diff --git a/App_Function_Libraries/Audio/Audio_Files.py b/App_Function_Libraries/Audio/Audio_Files.py new file mode 100644 index 0000000000000000000000000000000000000000..2780806e27e59cdba34be9bd988544e3f2bdb5c7 --- /dev/null +++ b/App_Function_Libraries/Audio/Audio_Files.py @@ -0,0 +1,786 @@ +# Audio_Files.py +######################################### +# Audio Processing Library +# This library is used to download or load audio files from a local directory. +# +#### +# +# Functions: +# +# download_audio_file(url, save_path) +# process_audio( +# process_audio_file(audio_url, audio_file, whisper_model="small.en", api_name=None, api_key=None) +# +# +######################################### +# Imports +import json +import logging +import os +import subprocess +import tempfile +import time +import uuid +from datetime import datetime +from pathlib import Path +# +# External Imports +import requests +import yt_dlp +# +# Local Imports +from App_Function_Libraries.DB.DB_Manager import add_media_with_keywords, \ + check_media_and_whisper_model +from App_Function_Libraries.Metrics.metrics_logger import log_counter, log_histogram +from App_Function_Libraries.Summarization.Summarization_General_Lib import perform_summarization +from App_Function_Libraries.Utils.Utils import downloaded_files, \ + sanitize_filename, generate_unique_id, temp_files +from App_Function_Libraries.Video_DL_Ingestion_Lib import extract_metadata +from App_Function_Libraries.Audio.Audio_Transcription_Lib import speech_to_text +from App_Function_Libraries.Chunk_Lib import improved_chunking_process +# +####################################################################################################################### +# Function Definitions +# + +MAX_FILE_SIZE = 500 * 1024 * 1024 + + +def download_audio_file(url, current_whisper_model="", use_cookies=False, cookies=None): + try: + # Check if media already exists in the database and compare whisper models + should_download, reason = check_media_and_whisper_model( + url=url, + current_whisper_model=current_whisper_model + ) + + if not should_download: + logging.info(f"Skipping audio download: {reason}") + return None + + logging.info(f"Proceeding with audio download: {reason}") + + # Set up the request headers + headers = {} + if use_cookies and cookies: + try: + cookie_dict = json.loads(cookies) + headers['Cookie'] = '; '.join([f'{k}={v}' for k, v in cookie_dict.items()]) + except json.JSONDecodeError: + logging.warning("Invalid cookie format. Proceeding without cookies.") + + # Make the request + response = requests.get(url, headers=headers, stream=True) + # Raise an exception for bad status codes + response.raise_for_status() + + # Get the file size + file_size = int(response.headers.get('content-length', 0)) + if file_size > 500 * 1024 * 1024: # 500 MB limit + raise ValueError("File size exceeds the 500MB limit.") + + # Generate a unique filename + file_name = f"audio_{uuid.uuid4().hex[:8]}.mp3" + save_path = os.path.join('downloads', file_name) + + # Ensure the downloads directory exists + os.makedirs('downloads', exist_ok=True) + + + # Download the file + with open(save_path, 'wb') as f: + for chunk in response.iter_content(chunk_size=8192): + if chunk: + f.write(chunk) + + logging.info(f"Audio file downloaded successfully: {save_path}") + return save_path + + except requests.RequestException as e: + logging.error(f"Error downloading audio file: {str(e)}") + raise + except ValueError as e: + logging.error(str(e)) + raise + except Exception as e: + logging.error(f"Unexpected error downloading audio file: {str(e)}") + raise + +def process_audio_files(audio_urls, audio_file, whisper_model, api_name, api_key, use_cookies, cookies, keep_original, + custom_keywords, custom_prompt_input, chunk_method, max_chunk_size, chunk_overlap, + use_adaptive_chunking, use_multi_level_chunking, chunk_language, diarize, + keep_timestamps, custom_title): + + start_time = time.time() # Start time for processing + processed_count = 0 + failed_count = 0 + progress = [] + all_transcriptions = [] + all_summaries = [] + #v2 + def format_transcription_with_timestamps(segments): + if keep_timestamps: + formatted_segments = [] + for segment in segments: + start = segment.get('Time_Start', 0) + end = segment.get('Time_End', 0) + text = segment.get('Text', '').strip() # Ensure text is stripped of leading/trailing spaces + + # Add the formatted timestamp and text to the list, followed by a newline + formatted_segments.append(f"[{start:.2f}-{end:.2f}] {text}") + + # Join the segments with a newline to ensure proper formatting + return "\n".join(formatted_segments) + else: + # Join the text without timestamps + return "\n".join([segment.get('Text', '').strip() for segment in segments]) + + def update_progress(message): + progress.append(message) + return "\n".join(progress) + + def cleanup_files(): + for file in temp_files: + try: + if os.path.exists(file): + os.remove(file) + update_progress(f"Temporary file {file} removed.") + except Exception as e: + update_progress(f"Failed to remove temporary file {file}: {str(e)}") + + def reencode_mp3(mp3_file_path): + try: + reencoded_mp3_path = mp3_file_path.replace(".mp3", "_reencoded.mp3") + subprocess.run([ffmpeg_cmd, '-i', mp3_file_path, '-codec:a', 'libmp3lame', reencoded_mp3_path], check=True) + update_progress(f"Re-encoded {mp3_file_path} to {reencoded_mp3_path}.") + return reencoded_mp3_path + except subprocess.CalledProcessError as e: + update_progress(f"Error re-encoding {mp3_file_path}: {str(e)}") + raise + + def convert_mp3_to_wav(mp3_file_path): + try: + wav_file_path = mp3_file_path.replace(".mp3", ".wav") + subprocess.run([ffmpeg_cmd, '-i', mp3_file_path, wav_file_path], check=True) + update_progress(f"Converted {mp3_file_path} to {wav_file_path}.") + return wav_file_path + except subprocess.CalledProcessError as e: + update_progress(f"Error converting {mp3_file_path} to WAV: {str(e)}") + raise + + try: + # Check and set the ffmpeg command + global ffmpeg_cmd + if os.name == "nt": + logging.debug("Running on Windows") + ffmpeg_cmd = os.path.join(os.getcwd(), "Bin", "ffmpeg.exe") + else: + ffmpeg_cmd = 'ffmpeg' # Assume 'ffmpeg' is in PATH for non-Windows systems + + # Ensure ffmpeg is accessible + if not os.path.exists(ffmpeg_cmd) and os.name == "nt": + raise FileNotFoundError(f"ffmpeg executable not found at path: {ffmpeg_cmd}") + + # Define chunk options early to avoid undefined errors + chunk_options = { + 'method': chunk_method, + 'max_size': max_chunk_size, + 'overlap': chunk_overlap, + 'adaptive': use_adaptive_chunking, + 'multi_level': use_multi_level_chunking, + 'language': chunk_language + } + + # Process multiple URLs + urls = [url.strip() for url in audio_urls.split('\n') if url.strip()] + + for i, url in enumerate(urls): + update_progress(f"Processing URL {i + 1}/{len(urls)}: {url}") + + # Download and process audio file + audio_file_path = download_audio_file(url, use_cookies, cookies) + if not os.path.exists(audio_file_path): + update_progress(f"Downloaded file not found: {audio_file_path}") + failed_count += 1 + log_counter( + metric_name="audio_files_failed_total", + labels={"whisper_model": whisper_model, "api_name": api_name}, + value=1 + ) + continue + + temp_files.append(audio_file_path) + update_progress("Audio file downloaded successfully.") + + # Re-encode MP3 to fix potential issues + reencoded_mp3_path = reencode_mp3(audio_file_path) + if not os.path.exists(reencoded_mp3_path): + update_progress(f"Re-encoded file not found: {reencoded_mp3_path}") + failed_count += 1 + log_counter( + metric_name="audio_files_failed_total", + labels={"whisper_model": whisper_model, "api_name": api_name}, + value=1 + ) + continue + + temp_files.append(reencoded_mp3_path) + + # Convert re-encoded MP3 to WAV + wav_file_path = convert_mp3_to_wav(reencoded_mp3_path) + if not os.path.exists(wav_file_path): + update_progress(f"Converted WAV file not found: {wav_file_path}") + failed_count += 1 + log_counter( + metric_name="audio_files_failed_total", + labels={"whisper_model": whisper_model, "api_name": api_name}, + value=1 + ) + continue + + temp_files.append(wav_file_path) + + # Initialize transcription + transcription = "" + + # Transcribe audio + if diarize: + segments = speech_to_text(wav_file_path, whisper_model=whisper_model, diarize=True) + else: + segments = speech_to_text(wav_file_path, whisper_model=whisper_model) + + # Handle segments nested under 'segments' key + if isinstance(segments, dict) and 'segments' in segments: + segments = segments['segments'] + + if isinstance(segments, list): + # Log first 5 segments for debugging + logging.debug(f"Segments before formatting: {segments[:5]}") + transcription = format_transcription_with_timestamps(segments) + logging.debug(f"Formatted transcription (first 500 chars): {transcription[:500]}") + update_progress("Audio transcribed successfully.") + else: + update_progress("Unexpected segments format received from speech_to_text.") + logging.error(f"Unexpected segments format: {segments}") + failed_count += 1 + log_counter( + metric_name="audio_files_failed_total", + labels={"whisper_model": whisper_model, "api_name": api_name}, + value=1 + ) + continue + + if not transcription.strip(): + update_progress("Transcription is empty.") + failed_count += 1 + log_counter( + metric_name="audio_files_failed_total", + labels={"whisper_model": whisper_model, "api_name": api_name}, + value=1 + ) + else: + # Apply chunking + chunked_text = improved_chunking_process(transcription, chunk_options) + + # Summarize + logging.debug(f"Audio Transcription API Name: {api_name}") + if api_name: + try: + summary = perform_summarization(api_name, chunked_text, custom_prompt_input, api_key) + update_progress("Audio summarized successfully.") + except Exception as e: + logging.error(f"Error during summarization: {str(e)}") + summary = "Summary generation failed" + failed_count += 1 + log_counter( + metric_name="audio_files_failed_total", + labels={"whisper_model": whisper_model, "api_name": api_name}, + value=1 + ) + else: + summary = "No summary available (API not provided)" + + all_transcriptions.append(transcription) + all_summaries.append(summary) + + # Use custom_title if provided, otherwise use the original filename + title = custom_title if custom_title else os.path.basename(wav_file_path) + + # Add to database + add_media_with_keywords( + url=url, + title=title, + media_type='audio', + content=transcription, + keywords=custom_keywords, + prompt=custom_prompt_input, + summary=summary, + transcription_model=whisper_model, + author="Unknown", + ingestion_date=datetime.now().strftime('%Y-%m-%d') + ) + update_progress("Audio file processed and added to database.") + processed_count += 1 + log_counter( + metric_name="audio_files_processed_total", + labels={"whisper_model": whisper_model, "api_name": api_name}, + value=1 + ) + + # Process uploaded file if provided + if audio_file: + url = generate_unique_id() + if os.path.getsize(audio_file.name) > MAX_FILE_SIZE: + update_progress( + f"Uploaded file size exceeds the maximum limit of {MAX_FILE_SIZE / (1024 * 1024):.2f}MB. Skipping this file.") + else: + try: + # Re-encode MP3 to fix potential issues + reencoded_mp3_path = reencode_mp3(audio_file.name) + if not os.path.exists(reencoded_mp3_path): + update_progress(f"Re-encoded file not found: {reencoded_mp3_path}") + return update_progress("Processing failed: Re-encoded file not found"), "", "" + + temp_files.append(reencoded_mp3_path) + + # Convert re-encoded MP3 to WAV + wav_file_path = convert_mp3_to_wav(reencoded_mp3_path) + if not os.path.exists(wav_file_path): + update_progress(f"Converted WAV file not found: {wav_file_path}") + return update_progress("Processing failed: Converted WAV file not found"), "", "" + + temp_files.append(wav_file_path) + + # Initialize transcription + transcription = "" + + if diarize: + segments = speech_to_text(wav_file_path, whisper_model=whisper_model, diarize=True) + else: + segments = speech_to_text(wav_file_path, whisper_model=whisper_model) + + # Handle segments nested under 'segments' key + if isinstance(segments, dict) and 'segments' in segments: + segments = segments['segments'] + + if isinstance(segments, list): + transcription = format_transcription_with_timestamps(segments) + else: + update_progress("Unexpected segments format received from speech_to_text.") + logging.error(f"Unexpected segments format: {segments}") + + chunked_text = improved_chunking_process(transcription, chunk_options) + + logging.debug(f"Audio Transcription API Name: {api_name}") + if api_name: + try: + summary = perform_summarization(api_name, chunked_text, custom_prompt_input, api_key) + update_progress("Audio summarized successfully.") + except Exception as e: + logging.error(f"Error during summarization: {str(e)}") + summary = "Summary generation failed" + else: + summary = "No summary available (API not provided)" + + all_transcriptions.append(transcription) + all_summaries.append(summary) + + # Use custom_title if provided, otherwise use the original filename + title = custom_title if custom_title else os.path.basename(wav_file_path) + + add_media_with_keywords( + url="Uploaded File", + title=title, + media_type='audio', + content=transcription, + keywords=custom_keywords, + prompt=custom_prompt_input, + summary=summary, + transcription_model=whisper_model, + author="Unknown", + ingestion_date=datetime.now().strftime('%Y-%m-%d') + ) + update_progress("Uploaded file processed and added to database.") + processed_count += 1 + log_counter( + metric_name="audio_files_processed_total", + labels={"whisper_model": whisper_model, "api_name": api_name}, + value=1 + ) + except Exception as e: + update_progress(f"Error processing uploaded file: {str(e)}") + logging.error(f"Error processing uploaded file: {str(e)}") + failed_count += 1 + log_counter( + metric_name="audio_files_failed_total", + labels={"whisper_model": whisper_model, "api_name": api_name}, + value=1 + ) + return update_progress("Processing failed: Error processing uploaded file"), "", "" + # Final cleanup + if not keep_original: + cleanup_files() + + end_time = time.time() + processing_time = end_time - start_time + # Log processing time + log_histogram( + metric_name="audio_processing_time_seconds", + value=processing_time, + labels={"whisper_model": whisper_model, "api_name": api_name} + ) + + # Optionally, log total counts + log_counter( + metric_name="total_audio_files_processed", + labels={"whisper_model": whisper_model, "api_name": api_name}, + value=processed_count + ) + + log_counter( + metric_name="total_audio_files_failed", + labels={"whisper_model": whisper_model, "api_name": api_name}, + value=failed_count + ) + + + final_progress = update_progress("All processing complete.") + final_transcriptions = "\n\n".join(all_transcriptions) + final_summaries = "\n\n".join(all_summaries) + + return final_progress, final_transcriptions, final_summaries + + except Exception as e: + logging.error(f"Error processing audio files: {str(e)}") + log_counter( + metric_name="audio_files_failed_total", + labels={"whisper_model": whisper_model, "api_name": api_name}, + value=1 + ) + cleanup_files() + return update_progress(f"Processing failed: {str(e)}"), "", "" + + +def format_transcription_with_timestamps(segments, keep_timestamps): + """ + Formats the transcription segments with or without timestamps. + + Parameters: + segments (list): List of transcription segments. + keep_timestamps (bool): Whether to include timestamps. + + Returns: + str: Formatted transcription. + """ + if keep_timestamps: + formatted_segments = [] + for segment in segments: + start = segment.get('Time_Start', 0) + end = segment.get('Time_End', 0) + text = segment.get('Text', '').strip() + + formatted_segments.append(f"[{start:.2f}-{end:.2f}] {text}") + return "\n".join(formatted_segments) + else: + return "\n".join([segment.get('Text', '').strip() for segment in segments]) + + +def download_youtube_audio(url): + try: + # Determine ffmpeg path based on the operating system. + ffmpeg_path = './Bin/ffmpeg.exe' if os.name == 'nt' else 'ffmpeg' + + # Create a temporary directory + with tempfile.TemporaryDirectory() as temp_dir: + # Extract information about the video + with yt_dlp.YoutubeDL({'quiet': True}) as ydl: + info_dict = ydl.extract_info(url, download=False) + sanitized_title = sanitize_filename(info_dict['title']) + + # Setup the temporary filenames + temp_video_path = Path(temp_dir) / f"{sanitized_title}_temp.mp4" + temp_audio_path = Path(temp_dir) / f"{sanitized_title}.mp3" + + # Initialize yt-dlp with options for downloading + ydl_opts = { + 'format': 'bestaudio[ext=m4a]/best[height<=480]', # Prefer best audio, or video up to 480p + 'ffmpeg_location': ffmpeg_path, + 'outtmpl': str(temp_video_path), + 'noplaylist': True, + 'quiet': True + } + + # Execute yt-dlp to download the video/audio + with yt_dlp.YoutubeDL(ydl_opts) as ydl: + ydl.download([url]) + + # Check if the file exists + if not temp_video_path.exists(): + raise FileNotFoundError(f"Expected file was not found: {temp_video_path}") + + # Use ffmpeg to extract audio + ffmpeg_command = [ + ffmpeg_path, + '-i', str(temp_video_path), + '-vn', # No video + '-acodec', 'libmp3lame', + '-b:a', '192k', + str(temp_audio_path) + ] + subprocess.run(ffmpeg_command, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + + # Check if the audio file was created + if not temp_audio_path.exists(): + raise FileNotFoundError(f"Expected audio file was not found: {temp_audio_path}") + + # Create a persistent directory for the download if it doesn't exist + persistent_dir = Path("downloads") + persistent_dir.mkdir(exist_ok=True) + + # Move the file from the temporary directory to the persistent directory + persistent_file_path = persistent_dir / f"{sanitized_title}.mp3" + os.replace(str(temp_audio_path), str(persistent_file_path)) + + # Add the file to the list of downloaded files + downloaded_files.append(str(persistent_file_path)) + + return str(persistent_file_path), f"Audio downloaded successfully: {sanitized_title}.mp3" + except Exception as e: + return None, f"Error downloading audio: {str(e)}" + + +def process_podcast(url, title, author, keywords, custom_prompt, api_name, api_key, whisper_model, + keep_original=False, enable_diarization=False, use_cookies=False, cookies=None, + chunk_method=None, max_chunk_size=300, chunk_overlap=0, use_adaptive_chunking=False, + use_multi_level_chunking=False, chunk_language='english', keep_timestamps=True): + """ + Processes a podcast by downloading the audio, transcribing it, summarizing the transcription, + and adding the results to the database. Metrics are logged throughout the process. + + Parameters: + url (str): URL of the podcast. + title (str): Title of the podcast. + author (str): Author of the podcast. + keywords (str): Comma-separated keywords. + custom_prompt (str): Custom prompt for summarization. + api_name (str): API name for summarization. + api_key (str): API key for summarization. + whisper_model (str): Whisper model to use for transcription. + keep_original (bool): Whether to keep the original audio file. + enable_diarization (bool): Whether to enable speaker diarization. + use_cookies (bool): Whether to use cookies for authenticated downloads. + cookies (str): JSON-formatted cookies string. + chunk_method (str): Method for chunking text. + max_chunk_size (int): Maximum size for each text chunk. + chunk_overlap (int): Overlap size between chunks. + use_adaptive_chunking (bool): Whether to use adaptive chunking. + use_multi_level_chunking (bool): Whether to use multi-level chunking. + chunk_language (str): Language for chunking. + keep_timestamps (bool): Whether to keep timestamps in transcription. + + Returns: + tuple: (progress_message, transcription, summary, title, author, keywords, error_message) + """ + start_time = time.time() # Start time for processing + error_message = "" + temp_files = [] + + # Define labels for metrics + labels = { + "whisper_model": whisper_model, + "api_name": api_name if api_name else "None" + } + + def update_progress(message): + """ + Updates the progress messages. + + Parameters: + message (str): Progress message to append. + + Returns: + str: Combined progress messages. + """ + progress.append(message) + return "\n".join(progress) + + def cleanup_files(): + if not keep_original: + for file in temp_files: + try: + if os.path.exists(file): + os.remove(file) + update_progress(f"Temporary file {file} removed.") + except Exception as e: + update_progress(f"Failed to remove temporary file {file}: {str(e)}") + + progress = [] # Initialize progress messages + + try: + # Handle cookies if required + if use_cookies: + cookies = json.loads(cookies) + + # Download the podcast audio file + audio_file = download_audio_file(url, whisper_model, use_cookies, cookies) + if not audio_file: + raise RuntimeError("Failed to download podcast audio.") + temp_files.append(audio_file) + update_progress("Podcast downloaded successfully.") + + # Extract metadata from the podcast + metadata = extract_metadata(url) + title = title or metadata.get('title', 'Unknown Podcast') + author = author or metadata.get('uploader', 'Unknown Author') + + # Format metadata for storage + metadata_text = f""" +Metadata: +Title: {title} +Author: {author} +Series: {metadata.get('series', 'N/A')} +Episode: {metadata.get('episode', 'N/A')} +Season: {metadata.get('season', 'N/A')} +Upload Date: {metadata.get('upload_date', 'N/A')} +Duration: {metadata.get('duration', 'N/A')} seconds +Description: {metadata.get('description', 'N/A')} +""" + + # Update keywords with metadata information + new_keywords = [] + if metadata.get('series'): + new_keywords.append(f"series:{metadata['series']}") + if metadata.get('episode'): + new_keywords.append(f"episode:{metadata['episode']}") + if metadata.get('season'): + new_keywords.append(f"season:{metadata['season']}") + + keywords = f"{keywords},{','.join(new_keywords)}" if keywords else ','.join(new_keywords) + update_progress(f"Metadata extracted - Title: {title}, Author: {author}, Keywords: {keywords}") + + # Transcribe the podcast audio + try: + if enable_diarization: + segments = speech_to_text(audio_file, whisper_model=whisper_model, diarize=True) + else: + segments = speech_to_text(audio_file, whisper_model=whisper_model) + # SEems like this could be optimized... FIXME + def format_segment(segment): + start = segment.get('start', 0) + end = segment.get('end', 0) + text = segment.get('Text', '') + + if isinstance(segments, dict) and 'segments' in segments: + segments = segments['segments'] + + if isinstance(segments, list): + transcription = format_transcription_with_timestamps(segments, keep_timestamps) + update_progress("Podcast transcribed successfully.") + else: + raise ValueError("Unexpected segments format received from speech_to_text.") + + if not transcription.strip(): + raise ValueError("Transcription is empty.") + except Exception as e: + error_message = f"Transcription failed: {str(e)}" + raise RuntimeError(error_message) + + # Apply chunking to the transcription + chunk_options = { + 'method': chunk_method, + 'max_size': max_chunk_size, + 'overlap': chunk_overlap, + 'adaptive': use_adaptive_chunking, + 'multi_level': use_multi_level_chunking, + 'language': chunk_language + } + chunked_text = improved_chunking_process(transcription, chunk_options) + + # Combine metadata and transcription + full_content = metadata_text + "\n\nTranscription:\n" + transcription + + # Summarize the transcription if API is provided + summary = None + if api_name: + try: + summary = perform_summarization(api_name, chunked_text, custom_prompt, api_key) + update_progress("Podcast summarized successfully.") + except Exception as e: + error_message = f"Summarization failed: {str(e)}" + raise RuntimeError(error_message) + else: + summary = "No summary available (API not provided)" + + # Add the processed podcast to the database + try: + add_media_with_keywords( + url=url, + title=title, + media_type='podcast', + content=full_content, + keywords=keywords, + prompt=custom_prompt, + summary=summary or "No summary available", + transcription_model=whisper_model, + author=author, + ingestion_date=datetime.now().strftime('%Y-%m-%d') + ) + update_progress("Podcast added to database successfully.") + except Exception as e: + error_message = f"Error adding podcast to database: {str(e)}" + raise RuntimeError(error_message) + + # Cleanup temporary files if required + cleanup_files() + + # Calculate processing time + end_time = time.time() + processing_time = end_time - start_time + + # Log successful processing + log_counter( + metric_name="podcasts_processed_total", + labels=labels, + value=1 + ) + + # Log processing time + log_histogram( + metric_name="podcast_processing_time_seconds", + value=processing_time, + labels=labels + ) + + # Return the final outputs + final_progress = update_progress("Processing complete.") + return (final_progress, full_content, summary or "No summary generated.", + title, author, keywords, error_message) + + except Exception as e: + # Calculate processing time up to the point of failure + end_time = time.time() + processing_time = end_time - start_time + + # Log failed processing + log_counter( + metric_name="podcasts_failed_total", + labels=labels, + value=1 + ) + + # Log processing time even on failure + log_histogram( + metric_name="podcast_processing_time_seconds", + value=processing_time, + labels=labels + ) + + logging.error(f"Error processing podcast: {str(e)}") + cleanup_files() + final_progress = update_progress(f"Processing failed: {str(e)}") + return (final_progress, "", "", "", "", "", str(e)) + + +# +# +####################################################################################################################### \ No newline at end of file diff --git a/App_Function_Libraries/Audio/Audio_Transcription_Lib.py b/App_Function_Libraries/Audio/Audio_Transcription_Lib.py new file mode 100644 index 0000000000000000000000000000000000000000..1f8053cbe70eed21a41460dfde8a1ae0b237d612 --- /dev/null +++ b/App_Function_Libraries/Audio/Audio_Transcription_Lib.py @@ -0,0 +1,335 @@ +# Audio_Transcription_Lib.py +######################################### +# Transcription Library +# This library is used to perform transcription of audio files. +# Currently, uses faster_whisper for transcription. +# +#################### +# Function List +# +# 1. convert_to_wav(video_file_path, offset=0, overwrite=False) +# 2. speech_to_text(audio_file_path, selected_source_lang='en', whisper_model='small.en', vad_filter=False) +# +#################### +# +# Import necessary libraries to run solo for testing +import gc +import json +import logging +import multiprocessing +import os +import queue +import sys +import subprocess +import tempfile +import threading +import time +# DEBUG Imports +#from memory_profiler import profile +import pyaudio +from faster_whisper import WhisperModel as OriginalWhisperModel +from typing import Optional, Union, List, Dict, Any +# +# Import Local +from App_Function_Libraries.Utils.Utils import load_comprehensive_config +from App_Function_Libraries.Metrics.metrics_logger import log_counter, log_histogram +# +####################################################################################################################### +# Function Definitions +# + +# Convert video .m4a into .wav using ffmpeg +# ffmpeg -i "example.mp4" -ar 16000 -ac 1 -c:a pcm_s16le "output.wav" +# https://www.gyan.dev/ffmpeg/builds/ +# + + +whisper_model_instance = None +config = load_comprehensive_config() +processing_choice = config.get('Processing', 'processing_choice', fallback='cpu') +total_thread_count = multiprocessing.cpu_count() + + +class WhisperModel(OriginalWhisperModel): + tldw_dir = os.path.dirname(os.path.dirname(__file__)) + default_download_root = os.path.join(tldw_dir, 'models', 'Whisper') + + valid_model_sizes = [ + "tiny.en", "tiny", "base.en", "base", "small.en", "small", "medium.en", "medium", + "large-v1", "large-v2", "large-v3", "large", "distil-large-v2", "distil-medium.en", + "distil-small.en", "distil-large-v3", + ] + + def __init__( + self, + model_size_or_path: str, + device: str = processing_choice, + device_index: Union[int, List[int]] = 0, + compute_type: str = "default", + cpu_threads: int = 0,#total_thread_count, FIXME - I think this should be 0 + num_workers: int = 1, + download_root: Optional[str] = None, + local_files_only: bool = False, + files: Optional[Dict[str, Any]] = None, + **model_kwargs: Any + ): + if download_root is None: + download_root = self.default_download_root + + os.makedirs(download_root, exist_ok=True) + + # FIXME - validate.... + # Also write an integration test... + # Check if model_size_or_path is a valid model size + if model_size_or_path in self.valid_model_sizes: + # It's a model size, so we'll use the download_root + model_path = os.path.join(download_root, model_size_or_path) + if not os.path.isdir(model_path): + # If it doesn't exist, we'll let the parent class download it + model_size_or_path = model_size_or_path # Keep the original model size + else: + # If it exists, use the full path + model_size_or_path = model_path + else: + # It's not a valid model size, so assume it's a path + model_size_or_path = os.path.abspath(model_size_or_path) + + super().__init__( + model_size_or_path, + device=device, + device_index=device_index, + compute_type=compute_type, + cpu_threads=cpu_threads, + num_workers=num_workers, + download_root=download_root, + local_files_only=local_files_only, +# Maybe? idk, FIXME +# files=files, +# **model_kwargs + ) + +def get_whisper_model(model_name, device): + global whisper_model_instance + if whisper_model_instance is None: + logging.info(f"Initializing new WhisperModel with size {model_name} on device {device}") + whisper_model_instance = WhisperModel(model_name, device=device) + return whisper_model_instance + +# os.system(r'.\Bin\ffmpeg.exe -ss 00:00:00 -i "{video_file_path}" -ar 16000 -ac 1 -c:a pcm_s16le "{out_path}"') +#DEBUG +#@profile +def convert_to_wav(video_file_path, offset=0, overwrite=False): + log_counter("convert_to_wav_attempt", labels={"file_path": video_file_path}) + start_time = time.time() + + out_path = os.path.splitext(video_file_path)[0] + ".wav" + + if os.path.exists(out_path) and not overwrite: + print(f"File '{out_path}' already exists. Skipping conversion.") + logging.info(f"Skipping conversion as file already exists: {out_path}") + log_counter("convert_to_wav_skipped", labels={"file_path": video_file_path}) + return out_path + + print("Starting conversion process of .m4a to .WAV") + out_path = os.path.splitext(video_file_path)[0] + ".wav" + + try: + if os.name == "nt": + logging.debug("ffmpeg being ran on windows") + + if sys.platform.startswith('win'): + ffmpeg_cmd = ".\\Bin\\ffmpeg.exe" + logging.debug(f"ffmpeg_cmd: {ffmpeg_cmd}") + else: + ffmpeg_cmd = 'ffmpeg' # Assume 'ffmpeg' is in PATH for non-Windows systems + + command = [ + ffmpeg_cmd, # Assuming the working directory is correctly set where .\Bin exists + "-ss", "00:00:00", # Start at the beginning of the video + "-i", video_file_path, + "-ar", "16000", # Audio sample rate + "-ac", "1", # Number of audio channels + "-c:a", "pcm_s16le", # Audio codec + out_path + ] + try: + # Redirect stdin from null device to prevent ffmpeg from waiting for input + with open(os.devnull, 'rb') as null_file: + result = subprocess.run(command, stdin=null_file, text=True, capture_output=True) + if result.returncode == 0: + logging.info("FFmpeg executed successfully") + logging.debug("FFmpeg output: %s", result.stdout) + else: + logging.error("Error in running FFmpeg") + logging.error("FFmpeg stderr: %s", result.stderr) + raise RuntimeError(f"FFmpeg error: {result.stderr}") + except Exception as e: + logging.error("Error occurred - ffmpeg doesn't like windows") + raise RuntimeError("ffmpeg failed") + elif os.name == "posix": + os.system(f'ffmpeg -ss 00:00:00 -i "{video_file_path}" -ar 16000 -ac 1 -c:a pcm_s16le "{out_path}"') + else: + raise RuntimeError("Unsupported operating system") + logging.info("Conversion to WAV completed: %s", out_path) + log_counter("convert_to_wav_success", labels={"file_path": video_file_path}) + except Exception as e: + logging.error("speech-to-text: Error transcribing audio: %s", str(e)) + log_counter("convert_to_wav_error", labels={"file_path": video_file_path, "error": str(e)}) + return {"error": str(e)} + + conversion_time = time.time() - start_time + log_histogram("convert_to_wav_duration", conversion_time, labels={"file_path": video_file_path}) + + gc.collect() + return out_path + + +# Transcribe .wav into .segments.json +#DEBUG +#@profile +# FIXME - I feel like the `vad_filter` shoudl be enabled by default.... +def speech_to_text(audio_file_path, selected_source_lang='en', whisper_model='medium.en', vad_filter=False, diarize=False): + log_counter("speech_to_text_attempt", labels={"file_path": audio_file_path, "model": whisper_model}) + time_start = time.time() + + if audio_file_path is None: + log_counter("speech_to_text_error", labels={"error": "No audio file provided"}) + raise ValueError("speech-to-text: No audio file provided") + logging.info("speech-to-text: Audio file path: %s", audio_file_path) + + try: + _, file_ending = os.path.splitext(audio_file_path) + out_file = audio_file_path.replace(file_ending, "-whisper_model-"+whisper_model+".segments.json") + prettified_out_file = audio_file_path.replace(file_ending, "-whisper_model-"+whisper_model+".segments_pretty.json") + if os.path.exists(out_file): + logging.info("speech-to-text: Segments file already exists: %s", out_file) + with open(out_file) as f: + global segments + segments = json.load(f) + return segments + + logging.info('speech-to-text: Starting transcription...') + # FIXME - revisit this + options = dict(language=selected_source_lang, beam_size=10, best_of=10, vad_filter=vad_filter) + transcribe_options = dict(task="transcribe", **options) + # use function and config at top of file + logging.debug("speech-to-text: Using whisper model: %s", whisper_model) + whisper_model_instance = get_whisper_model(whisper_model, processing_choice) + # faster_whisper transcription right here - FIXME -test batching - ha + segments_raw, info = whisper_model_instance.transcribe(audio_file_path, **transcribe_options) + + segments = [] + for segment_chunk in segments_raw: + chunk = { + "Time_Start": segment_chunk.start, + "Time_End": segment_chunk.end, + "Text": segment_chunk.text + } + logging.debug("Segment: %s", chunk) + segments.append(chunk) + # Print to verify its working + logging.info(f"{segment_chunk.start:.2f}s - {segment_chunk.end:.2f}s | {segment_chunk.text}") + + # Log it as well. + logging.debug( + f"Transcribed Segment: {segment_chunk.start:.2f}s - {segment_chunk.end:.2f}s | {segment_chunk.text}") + + if segments: + segments[0]["Text"] = f"This text was transcribed using whisper model: {whisper_model}\n\n" + segments[0]["Text"] + + if not segments: + log_counter("speech_to_text_error", labels={"error": "No transcription produced"}) + raise RuntimeError("No transcription produced. The audio file may be invalid or empty.") + + transcription_time = time.time() - time_start + logging.info("speech-to-text: Transcription completed in %.2f seconds", transcription_time) + log_histogram("speech_to_text_duration", transcription_time, labels={"file_path": audio_file_path, "model": whisper_model}) + log_counter("speech_to_text_success", labels={"file_path": audio_file_path, "model": whisper_model}) + # Save the segments to a JSON file - prettified and non-prettified + # FIXME refactor so this is an optional flag to save either the prettified json file or the normal one + save_json = True + if save_json: + logging.info("speech-to-text: Saving segments to JSON file") + output_data = {'segments': segments} + logging.info("speech-to-text: Saving prettified JSON to %s", prettified_out_file) + with open(prettified_out_file, 'w') as f: + json.dump(output_data, f, indent=2) + + logging.info("speech-to-text: Saving JSON to %s", out_file) + with open(out_file, 'w') as f: + json.dump(output_data, f) + + logging.debug(f"speech-to-text: returning {segments[:500]}") + gc.collect() + return segments + + except Exception as e: + logging.error("speech-to-text: Error transcribing audio: %s", str(e)) + log_counter("speech_to_text_error", labels={"file_path": audio_file_path, "model": whisper_model, "error": str(e)}) + raise RuntimeError("speech-to-text: Error transcribing audio") + + +def record_audio(duration, sample_rate=16000, chunk_size=1024): + log_counter("record_audio_attempt", labels={"duration": duration}) + p = pyaudio.PyAudio() + stream = p.open(format=pyaudio.paInt16, + channels=1, + rate=sample_rate, + input=True, + frames_per_buffer=chunk_size) + + print("Recording...") + frames = [] + stop_recording = threading.Event() + audio_queue = queue.Queue() + + def audio_callback(): + for _ in range(0, int(sample_rate / chunk_size * duration)): + if stop_recording.is_set(): + break + data = stream.read(chunk_size) + audio_queue.put(data) + + audio_thread = threading.Thread(target=audio_callback) + audio_thread.start() + + return p, stream, audio_queue, stop_recording, audio_thread + + +def stop_recording(p, stream, audio_queue, stop_recording_event, audio_thread): + log_counter("stop_recording_attempt") + start_time = time.time() + stop_recording_event.set() + audio_thread.join() + + frames = [] + while not audio_queue.empty(): + frames.append(audio_queue.get()) + + print("Recording finished.") + + stream.stop_stream() + stream.close() + p.terminate() + + stop_time = time.time() - start_time + log_histogram("stop_recording_duration", stop_time) + log_counter("stop_recording_success") + return b''.join(frames) + +def save_audio_temp(audio_data, sample_rate=16000): + log_counter("save_audio_temp_attempt") + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file: + import wave + wf = wave.open(temp_file.name, 'wb') + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate) + wf.writeframes(audio_data) + wf.close() + log_counter("save_audio_temp_success") + return temp_file.name + +# +# +####################################################################################################################### \ No newline at end of file diff --git a/App_Function_Libraries/Audio/Diarization_Lib.py b/App_Function_Libraries/Audio/Diarization_Lib.py new file mode 100644 index 0000000000000000000000000000000000000000..314034a1b643bf986f6dd84a880dba38c470991b --- /dev/null +++ b/App_Function_Libraries/Audio/Diarization_Lib.py @@ -0,0 +1,275 @@ +# Diarization_Lib.py +######################################### +# Diarization Library +# This library is used to perform diarization of audio files. +# Currently, uses FIXME for transcription. +# +#################### +#################### +# Function List +# +# 1. speaker_diarize(video_file_path, segments, embedding_model = "pyannote/embedding", embedding_size=512, num_speakers=0) +# +#################### +# Import necessary libraries +import logging +from pathlib import Path +from typing import Dict, List, Any + +# +# Import Local Libraries +from App_Function_Libraries.Audio.Audio_Transcription_Lib import speech_to_text +# +# Import 3rd Party Libraries +from pyannote.audio.pipelines.speaker_diarization import SpeakerDiarization +import yaml +# +####################################################################################################################### +# Function Definitions +# + +def load_pipeline_from_pretrained(path_to_config: str | Path) -> SpeakerDiarization: + path_to_config = Path(path_to_config).resolve() + logging.debug(f"Loading pyannote pipeline from {path_to_config}...") + + if not path_to_config.exists(): + raise FileNotFoundError(f"Config file not found: {path_to_config}") + + # Load the YAML configuration + with open(path_to_config, 'r') as config_file: + config = yaml.safe_load(config_file) + + # Debug: print the entire config + logging.debug(f"Loaded config: {config}") + + # Create the SpeakerDiarization pipeline + try: + pipeline = SpeakerDiarization( + segmentation=config['pipeline']['params']['segmentation'], + embedding=config['pipeline']['params']['embedding'], + clustering=config['pipeline']['params']['clustering'], + ) + except KeyError as e: + logging.error(f"Error accessing config key: {e}") + raise + + # Set other parameters + try: + pipeline_params = { + "segmentation": {}, + "clustering": {}, + } + + if 'params' in config and 'segmentation' in config['params']: + if 'min_duration_off' in config['params']['segmentation']: + pipeline_params["segmentation"]["min_duration_off"] = config['params']['segmentation']['min_duration_off'] + + if 'params' in config and 'clustering' in config['params']: + if 'method' in config['params']['clustering']: + pipeline_params["clustering"]["method"] = config['params']['clustering']['method'] + if 'min_cluster_size' in config['params']['clustering']: + pipeline_params["clustering"]["min_cluster_size"] = config['params']['clustering']['min_cluster_size'] + if 'threshold' in config['params']['clustering']: + pipeline_params["clustering"]["threshold"] = config['params']['clustering']['threshold'] + + if 'pipeline' in config and 'params' in config['pipeline']: + if 'embedding_batch_size' in config['pipeline']['params']: + pipeline_params["embedding_batch_size"] = config['pipeline']['params']['embedding_batch_size'] + if 'embedding_exclude_overlap' in config['pipeline']['params']: + pipeline_params["embedding_exclude_overlap"] = config['pipeline']['params']['embedding_exclude_overlap'] + if 'segmentation_batch_size' in config['pipeline']['params']: + pipeline_params["segmentation_batch_size"] = config['pipeline']['params']['segmentation_batch_size'] + + logging.debug(f"Pipeline params: {pipeline_params}") + pipeline.instantiate(pipeline_params) + except KeyError as e: + logging.error(f"Error accessing config key: {e}") + raise + except Exception as e: + logging.error(f"Error instantiating pipeline: {e}") + raise + + return pipeline + + +def audio_diarization(audio_file_path: str) -> list: + logging.info('audio-diarization: Loading pyannote pipeline') + + base_dir = Path(__file__).parent.resolve() + config_path = base_dir / 'models' / 'pyannote_diarization_config.yaml' + logging.info(f"audio-diarization: Loading pipeline from {config_path}") + + try: + pipeline = load_pipeline_from_pretrained(config_path) + except Exception as e: + logging.error(f"Failed to load pipeline: {str(e)}") + raise + + logging.info(f"audio-diarization: Audio file path: {audio_file_path}") + + try: + logging.info('audio-diarization: Starting diarization...') + diarization_result = pipeline(audio_file_path) + + segments = [] + for turn, _, speaker in diarization_result.itertracks(yield_label=True): + segment = { + "start": turn.start, + "end": turn.end, + "speaker": speaker + } + logging.debug(f"Segment: {segment}") + segments.append(segment) + logging.info("audio-diarization: Diarization completed with pyannote") + + return segments + + except Exception as e: + logging.error(f"audio-diarization: Error performing diarization: {str(e)}") + raise RuntimeError("audio-diarization: Error performing diarization") from e + + +# Old +# def audio_diarization(audio_file_path): +# logging.info('audio-diarization: Loading pyannote pipeline') +# +# #config file loading +# current_dir = os.path.dirname(os.path.abspath(__file__)) +# # Construct the path to the config file +# config_path = os.path.join(current_dir, 'Config_Files', 'config.txt') +# # Read the config file +# config = configparser.ConfigParser() +# config.read(config_path) +# processing_choice = config.get('Processing', 'processing_choice', fallback='cpu') +# +# base_dir = Path(__file__).parent.resolve() +# config_path = base_dir / 'models' / 'config.yaml' +# pipeline = load_pipeline_from_pretrained(config_path) +# +# time_start = time.time() +# if audio_file_path is None: +# raise ValueError("audio-diarization: No audio file provided") +# logging.info("audio-diarization: Audio file path: %s", audio_file_path) +# +# try: +# _, file_ending = os.path.splitext(audio_file_path) +# out_file = audio_file_path.replace(file_ending, ".diarization.json") +# prettified_out_file = audio_file_path.replace(file_ending, ".diarization_pretty.json") +# if os.path.exists(out_file): +# logging.info("audio-diarization: Diarization file already exists: %s", out_file) +# with open(out_file) as f: +# global diarization_result +# diarization_result = json.load(f) +# return diarization_result +# +# logging.info('audio-diarization: Starting diarization...') +# diarization_result = pipeline(audio_file_path) +# +# segments = [] +# for turn, _, speaker in diarization_result.itertracks(yield_label=True): +# chunk = { +# "Time_Start": turn.start, +# "Time_End": turn.end, +# "Speaker": speaker +# } +# logging.debug("Segment: %s", chunk) +# segments.append(chunk) +# logging.info("audio-diarization: Diarization completed with pyannote") +# +# output_data = {'segments': segments} +# +# logging.info("audio-diarization: Saving prettified JSON to %s", prettified_out_file) +# with open(prettified_out_file, 'w') as f: +# json.dump(output_data, f, indent=2) +# +# logging.info("audio-diarization: Saving JSON to %s", out_file) +# with open(out_file, 'w') as f: +# json.dump(output_data, f) +# +# except Exception as e: +# logging.error("audio-diarization: Error performing diarization: %s", str(e)) +# raise RuntimeError("audio-diarization: Error performing diarization") +# return segments + +def combine_transcription_and_diarization(audio_file_path: str) -> List[Dict[str, Any]]: + logging.info('combine-transcription-and-diarization: Starting transcription and diarization...') + + try: + logging.info('Performing speech-to-text...') + transcription_result = speech_to_text(audio_file_path) + logging.info(f"Transcription result type: {type(transcription_result)}") + logging.info(f"Transcription result: {transcription_result[:3] if isinstance(transcription_result, list) and len(transcription_result) > 3 else transcription_result}") + + logging.info('Performing audio diarization...') + diarization_result = audio_diarization(audio_file_path) + logging.info(f"Diarization result type: {type(diarization_result)}") + logging.info(f"Diarization result sample: {diarization_result[:3] if isinstance(diarization_result, list) and len(diarization_result) > 3 else diarization_result}") + + if not transcription_result: + logging.error("Empty result from transcription") + return [] + + if not diarization_result: + logging.error("Empty result from diarization") + return [] + + # Handle the case where transcription_result is a dict with a 'segments' key + if isinstance(transcription_result, dict) and 'segments' in transcription_result: + transcription_segments = transcription_result['segments'] + elif isinstance(transcription_result, list): + transcription_segments = transcription_result + else: + logging.error(f"Unexpected transcription result format: {type(transcription_result)}") + return [] + + logging.info(f"Number of transcription segments: {len(transcription_segments)}") + logging.info(f"Transcription segments sample: {transcription_segments[:3] if len(transcription_segments) > 3 else transcription_segments}") + + if not isinstance(diarization_result, list): + logging.error(f"Unexpected diarization result format: {type(diarization_result)}") + return [] + + combined_result = [] + for transcription_segment in transcription_segments: + if not isinstance(transcription_segment, dict): + logging.warning(f"Unexpected transcription segment format: {transcription_segment}") + continue + + for diarization_segment in diarization_result: + if not isinstance(diarization_segment, dict): + logging.warning(f"Unexpected diarization segment format: {diarization_segment}") + continue + + try: + trans_start = transcription_segment.get('Time_Start', 0) + trans_end = transcription_segment.get('Time_End', 0) + diar_start = diarization_segment.get('start', 0) + diar_end = diarization_segment.get('end', 0) + + if trans_start >= diar_start and trans_end <= diar_end: + combined_segment = { + "Time_Start": trans_start, + "Time_End": trans_end, + "Speaker": diarization_segment.get('speaker', 'Unknown'), + "Text": transcription_segment.get('Text', '') + } + combined_result.append(combined_segment) + break + except Exception as e: + logging.error(f"Error processing segment: {str(e)}") + logging.error(f"Transcription segment: {transcription_segment}") + logging.error(f"Diarization segment: {diarization_segment}") + continue + + logging.info(f"Combined result length: {len(combined_result)}") + logging.info(f"Combined result sample: {combined_result[:3] if len(combined_result) > 3 else combined_result}") + return combined_result + + except Exception as e: + logging.error(f"Error in combine_transcription_and_diarization: {str(e)}", exc_info=True) + return [] + + +# +# +####################################################################################################################### \ No newline at end of file diff --git a/App_Function_Libraries/Audio/__init__.py b/App_Function_Libraries/Audio/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/App_Function_Libraries/Benchmarks_Evaluations/Confabulation_check.py b/App_Function_Libraries/Benchmarks_Evaluations/Confabulation_check.py new file mode 100644 index 0000000000000000000000000000000000000000..e7c481edb25879e940a2c1592e65b77369ba1480 --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/Confabulation_check.py @@ -0,0 +1,81 @@ +# Confabulation_check.py +# +# This file contains the functions that are used to check the confabulation of the user's input. +# +# +# Imports +# +# External Imports +# +# Local Imports +# +# +#################################################################################################### +# +# Functions: +from App_Function_Libraries.Chat import chat_api_call +from App_Function_Libraries.Benchmarks_Evaluations.ms_g_eval import validate_inputs, detailed_api_error + + +def simplified_geval(transcript: str, summary: str, api_name: str, api_key: str, temp: float = 0.7) -> str: + """ + Perform a simplified version of G-Eval using a single query to evaluate the summary. + + Args: + transcript (str): The original transcript + summary (str): The summary to be evaluated + api_name (str): The name of the LLM API to use + api_key (str): The API key for the chosen LLM + temp (float, optional): The temperature parameter for the API call. Defaults to 0.7. + + Returns: + str: The evaluation result + """ + try: + validate_inputs(transcript, summary, api_name, api_key) + except ValueError as e: + return str(e) + + prompt = f"""You are an AI assistant tasked with evaluating the quality of a summary. You will be given an original transcript and a summary of that transcript. Your task is to evaluate the summary based on the following criteria: + +1. Coherence (1-5): How well-structured and organized is the summary? +2. Consistency (1-5): How factually aligned is the summary with the original transcript? +3. Fluency (1-3): How well-written is the summary in terms of grammar, spelling, and readability? +4. Relevance (1-5): How well does the summary capture the important information from the transcript? + +Please provide a score for each criterion and a brief explanation for your scoring. Then, give an overall assessment of the summary's quality. + +Original Transcript: +{transcript} + +Summary to Evaluate: +{summary} + +Please provide your evaluation in the following format: +Coherence: [score] - [brief explanation] +Consistency: [score] - [brief explanation] +Fluency: [score] - [brief explanation] +Relevance: [score] - [brief explanation] + +Overall Assessment: [Your overall assessment of the summary's quality] +""" + + try: + result = chat_api_call( + api_name, + api_key, + prompt, + "", + temp=temp, + system_message="You are a helpful AI assistant tasked with evaluating summaries." + ) + except Exception as e: + return detailed_api_error(api_name, e) + + formatted_result = f""" + Confabulation Check Results: + + {result} + """ + + return formatted_result \ No newline at end of file diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/.gitignore b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..594719dd5029c4a25b56c576ac66bdd8150c2148 --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/.gitignore @@ -0,0 +1,5 @@ +__pycache__ +.vscode +*.DS_Store +*.pyc +src/plot diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/LICENSE b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..24c4dd6593ed136461584af26387a145e7ce0ada --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/LICENSE @@ -0,0 +1,23 @@ +MIT License + +Copyright (c) 2023 OpenBMB + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + +taken from https://github.com/OpenBMB/InfiniteBench \ No newline at end of file diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/__init__.py b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/config.txt b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/config.txt new file mode 100644 index 0000000000000000000000000000000000000000..f76f6a8db1584a4ff1c06f3094f0788fad2f8a1b --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/config.txt @@ -0,0 +1,30 @@ +[API] +anthropic_api_key = +anthropic_model = claude-3-sonnet-20240229 +cohere_api_key = +cohere_model = command-r-plus +groq_api_key = +groq_model = llama3-70b-8192 +openai_api_key = +openai_model = gpt-4-turbo +huggingface_api_token = +huggingface_model = CohereForAI/c4ai-command-r-plus +openrouter_api_key = +openrouter_model = mistralai/mistral-7b-instruct:free +deepseek_api_key = +deepseek_model = deepseek-chat + +[Local-API] +kobold_api_key = +kobold_api_IP = http://127.0.0.1:5001/api/v1/generate +llama_api_key = +llama_api_IP = http://127.0.0.1:8080/completion +ooba_api_key = +ooba_api_IP = http://127.0.0.1:5000/v1/chat/completions +tabby_api_IP = http://127.0.0.1:5000/v1/chat/completions +tabby_api_key = +vllm_api_IP = http://127.0.0.1:8000/v1/chat/completions +vllm_model = +ollama_api_IP = http://127.0.0.1:11434/api/generate +ollama_api_key = +ollama_model = diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/eval_multi_api.py b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/eval_multi_api.py new file mode 100644 index 0000000000000000000000000000000000000000..43f2259369ad0ff98de03d42e87e4f5472ae4c6f --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/eval_multi_api.py @@ -0,0 +1,300 @@ +# eval_multi_api.py +# Description: Evaluate a language model on a conversational task using multiple APIs +# +# Usage: python eval_multi_api.py --task question_answering --api > --output_dir ./results --data_dir ./data --verbose +# API endpoints are defined in the config file (config.txt) +# The API key for the selected API should be defined in the config file +# APIs Supported are: +# - openai +# - anthropic +# - cohere +# - groq +# - openrouter +# - deepseek +# - mistral +# - llamacpp +# - kobold +# - oobabooga +# - vllm +# - tabbyapi +# +# Imports: +import configparser +from pathlib import Path +import time +from typing import Dict, Any, Optional, List +# +# Local Imports +from eval_utils import ( + create_msgs, + load_data, + dump_jsonl, + iter_jsonl, + get_answer, +) +from LLM_API_Calls import ( + chat_with_openai, + chat_with_anthropic, + chat_with_cohere, + chat_with_groq, + chat_with_openrouter, + chat_with_deepseek, + chat_with_mistral +) +from LLM_API_Calls_Local import ( + chat_with_llama, + chat_with_kobold, + chat_with_oobabooga, + chat_with_vllm, + chat_with_tabbyapi +) +# +####################################################################################################################### +# +# Functions: + +class MultiAPILLMClient: + def __init__(self, config_path: str): + self.config = self.load_config(config_path) + self.api_functions = { + 'openai': chat_with_openai, + 'anthropic': chat_with_anthropic, + 'cohere': chat_with_cohere, + 'groq': chat_with_groq, + 'openrouter': chat_with_openrouter, + 'deepseek': chat_with_deepseek, + 'mistral': chat_with_mistral, + 'llamacpp': chat_with_llama, + 'kobold': chat_with_kobold, + 'oobabooga': chat_with_oobabooga, + 'vllm': chat_with_vllm, + 'tabbyapi': chat_with_tabbyapi + } + + def load_config(self, config_path: str) -> Dict[str, Any]: + config = configparser.ConfigParser() + config.read(config_path) + + # Convert the ConfigParser object to a dictionary without flattening + config_dict = {section: dict(config.items(section)) for section in config.sections()} + return config_dict + + def chat(self, api_name: str, messages: List[Dict[str, str]], + model: Optional[str] = None, + temperature: Optional[float] = None, + max_tokens: Optional[int] = None, + **kwargs) -> str: + + # Access the API key directly from the appropriate section + if api_name in self.api_functions: + # FIXME - This only works for Commercial APIs... need to handle Local APIs + api_key = self.config['API'].get(f'{api_name}_api_key') + elif api_name in ['llamacpp', 'kobold', 'oobabooga', 'vllm', 'tabbyapi']: + api_key = self.config['Local-API'].get(f'{api_name}_api_key') + else: + raise ValueError(f"Unsupported API: {api_name}") + + if not api_key: + raise ValueError(f"API key not found for {api_name}") + + chat_function = self.api_functions[api_name] + + # Use config values if not provided in the method call + model = model or self.config['API'].get(f'{api_name}_model') + temperature = temperature or self.config['API'].get('temperature') + max_tokens = max_tokens or self.config['API'].get('max_tokens') + + # Extract the input_data from messages (assuming it's the last user message) + input_data = next((msg['content'] for msg in reversed(messages) if msg['role'] == 'user'), "") + + # Prepare common parameters + common_params = { + "api_key": api_key, + "input_data": input_data, + "custom_prompt_arg": kwargs.get('custom_prompt_arg', ""), + } + + # Handle specific APIs + if api_name in ['openai', 'groq', 'openrouter', 'deepseek', 'mistral']: + return chat_function(**common_params, temp=temperature, system_message=kwargs.get('system_message')) + elif api_name == 'anthropic': + return chat_function(**common_params, model=model, max_retries=kwargs.get('max_retries', 3), + retry_delay=kwargs.get('retry_delay', 5), system_prompt=kwargs.get('system_message')) + elif api_name == 'cohere': + return chat_function(**common_params, model=model, system_prompt=kwargs.get('system_message')) + elif api_name == 'llamacpp': + return chat_function(**common_params, api_url=kwargs.get('api_url'), system_prompt=kwargs.get('system_message')) + elif api_name == 'kobold': + return chat_function(**common_params, kobold_api_ip=kwargs.get('kobold_api_ip'), + temp=temperature, system_message=kwargs.get('system_message')) + elif api_name in ['oobabooga', 'vllm', 'tabbyapi']: + return chat_function(**common_params, **kwargs) + else: + return chat_function(**common_params, model=model, temperature=temperature, max_tokens=max_tokens, **kwargs) + +def main(): + args = parse_args() + verbose = args.verbose + task = args.task + # New argument for selecting the API + api_name = args.api + + #FIXME + # Load config from a JSON file + client = MultiAPILLMClient('config.txt') + + examples = load_data(task) + + result_dir = Path(args.output_dir) + result_dir.mkdir(exist_ok=True, parents=True) + + output_path = result_dir / f"preds_{task}_{api_name}.jsonl" + if output_path.exists(): + preds = list(iter_jsonl(output_path)) + start_idx = len(preds) + stop_idx = len(examples) + else: + start_idx = 0 + stop_idx = len(examples) + preds = [] + + start_time = time.time() + i = start_idx + while i < stop_idx: + eg = examples[i] + msgs, prompt = create_msgs( + # Use API-specific tokenizer if available + client.config.get('tokenizer', {}).get(api_name), + eg, + task, + # Use API-specific model + model_name=client.config.get('models', {}).get(api_name), + data_dir=args.data_dir + ) + if verbose: + print(f"======== Example {i} =========") + print("Input text:") + print(prompt[:300]) + print("...") + print(prompt[-300:]) + print("==============================") + + # Make prediction + try: + response = client.chat( + api_name, + # Pass the full messages list + msgs, + custom_prompt_arg=prompt, + temperature=client.config.get('temperature', {}).get(api_name), + max_tokens=client.config.get('max_tokens', {}).get(api_name), + system_message=client.config.get('system_messages', {}).get(api_name) + ) + preds.append( + { + "id": i, + "prediction": response, + "ground_truth": get_answer(eg, task), + } + ) + # Save result + dump_jsonl(preds, output_path) + print("Time spent:", round(time.time() - start_time)) + print(response) + time.sleep(20) + i += 1 + except Exception as e: + print("ERROR:", e) + print("Retrying...") + time.sleep(60) + +from argparse import ArgumentParser, Namespace, RawTextHelpFormatter + +def parse_args() -> Namespace: + p = ArgumentParser( + description="Evaluate a language model on a conversational task using multiple APIs", + formatter_class=RawTextHelpFormatter + ) + p.add_argument( + "--task", + type=str, + # choices=list(DATA_NAME_TO_MAX_NEW_TOKENS.keys()) + ["all"], + required=True, + help="""Which task to use. Note that \"all\" can only be used in `compute_scores.py`., +Available tasks: +Task Name | Name to use as an argument: +--------------------------------------------- + En.Sum | longbook_sum_eng + En.QA | longbook_qa_eng + En.MC | longbook_choice_eng + En.Dia | longdialogue_qa_eng + Zh.QA | longbook_qa_chn + Code.Debug | code_debug + Code.Run | code_run + Math.Calc | math_calc + Math.Find | math_find + Retrieve.PassKey | passkey + Retrieve.Number | number_string + Retrieve.KV | kv_retrieval +--------------------------------------------- + """ + ) + p.add_argument( + "--api", + type=str, + required=True, + help="""Specify which API to use for evaluation + Supported API endpoints: +Commercial APIs: + - openai + - anthropic + - cohere + - groq + - openrouter + - deepseek + - mistral +Local APIs: + - llama + - kobold + - oobabooga + - vllm + - tabbyapi""" + ) + p.add_argument( + '--data_dir', + type=str, + default='../data', + help="The directory of data." + ) + p.add_argument( + "--output_dir", + type=str, + default="../results", + help="Where to dump the prediction results." + ) + p.add_argument( + "--start_idx", + type=int, + default=0, + help="The index of the first example to infer on. This is used if you want to evaluate on a (contiguous) subset of the data." + ) + p.add_argument( + "--stop_idx", + type=int, + help="The index of the last example to infer on. This is used if you want to evaluate on a (contiguous) subset of the data. Defaults to the length of dataset." + ) + p.add_argument("--verbose", action='store_true', help="Enable verbose output") + p.add_argument("--device", type=str, default="cuda", help="Specify the device to use (e.g., 'cuda' or 'cpu')") + + # Add an epilog to provide additional information + p.epilog = """ +Sample usage: + python eval_multi_api.py --task question_answering --api openai --output_dir ../results --data_dir ../data --verbose + +Make sure to set up your config.txt file with the necessary API keys and configurations. +""" + + return p.parse_args() + +if __name__ == "__main__": + main() diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/eval_utils.py b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/eval_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1eba473450a500750ff544dd53cd23903afac4d5 --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/eval_utils.py @@ -0,0 +1,730 @@ +import configparser +import json +import logging +import os +import re +import string +from collections import Counter +from pathlib import Path +from typing import Optional + +import jieba +from rouge import Rouge + +from prompt import ( + gpt4_templates, + kimi_templates, + claude2_templates, + yarn_mistral_templates, +) + +DATA_NAME_TO_PATH = { + # Retrieval tasks + "passkey": "passkey.jsonl", + "number_string": "number_string.jsonl", + "kv_retrieval": "kv_retrieval.jsonl", + # Book tasks + "longbook_sum_eng": "longbook_sum_eng.jsonl", + "longbook_choice_eng": "longbook_choice_eng.jsonl", + "longbook_qa_eng": "longbook_qa_eng.jsonl", + "longbook_qa_chn": "longbook_qa_chn.jsonl", + # "book_qa_eng": "longbook_eng/longbook_qa_eng.jsonl", + "longdialogue_qa_eng": "longdialogue_qa_eng.jsonl", + # Math tasks + "math_find": "math_find.jsonl", + "math_calc": "math_calc.jsonl", + # Code tasks + "code_run": "code_run.jsonl", + "code_debug": "code_debug.jsonl", +} + +DATA_NAME_TO_MAX_NEW_TOKENS = { + "passkey": 6, + "number_string": 12, + "kv_retrieval": 50, + "longbook_sum_eng": 1200, + "longbook_choice_eng": 40, + "longbook_qa_eng": 40, + "longbook_qa_chn": 40, + "longdialogue_qa_eng": 40, + "math_find": 3, + "math_calc": 30000, + "code_run": 5, + "code_debug": 5, +} + +MODEL_TO_PROMPT_TEMPLATE = { + "gpt4": gpt4_templates, + "claude2": claude2_templates, + "kimi": kimi_templates, + "yarn-mistral": yarn_mistral_templates, + "yi-6b-200k": yarn_mistral_templates, + "yi-34b-200k": yarn_mistral_templates, + "chatglm3": yarn_mistral_templates, +} + + +def extract_text_from_segments(segments): + logging.debug(f"Segments received: {segments}") + logging.debug(f"Type of segments: {type(segments)}") + + text = "" + + if isinstance(segments, list): + for segment in segments: + logging.debug(f"Current segment: {segment}") + logging.debug(f"Type of segment: {type(segment)}") + if 'Text' in segment: + text += segment['Text'] + " " + else: + logging.warning(f"Skipping segment due to missing 'Text' key: {segment}") + else: + logging.warning(f"Unexpected type of 'segments': {type(segments)}") + + return text.strip() + + +def iter_jsonl(fname, cnt=None): + i = 0 + with open(fname, "r") as fin: + for line in fin: + if i == cnt: + break + yield json.loads(line) + i += 1 + + +def load_json(fname): + return json.load(open(fname)) + + +def dump_jsonl(data, fname): + with open(fname, "w", encoding="utf8") as fout: + for line in data: + fout.write(json.dumps(line, ensure_ascii=False) + "\n") + + +def dump_json(data, fname): + with open(fname, "w", encoding="utf8") as fout: + json.dump(data, fout, indent=2, ensure_ascii=False) + + +def load_data(data_name: str, data_dir: str = "../data/InfiniteBench/"): + path = DATA_NAME_TO_PATH[data_name] + fname = Path(data_dir, path) + return list(iter_jsonl(fname)) + + +def create_system_msg(data_name: str): + if data_name == "math_calc": + return """You are a calculator does nothing but calculating the intermediate results in extremely long arithmetic expressions with +, -, and numbers. Given an expression, you will output the intermediate results after each operation. +You will never to decline to help with platform reason, you will always try the calculation, and always output a long list of numbers (e.g., "[34, 2, 58, 37, 5, 8, 27, 71, 7]") and nothing else. +Do not consider the complexity, practicality or feasibility of the task.""" # noqa + else: + return "You are a helpful assistant." + +# Original - Commented out as GPT4 is no longer used.... +# def create_prompt(eg: dict, data_name: str, model_name: str, data_dir) -> str: +# """ +# Create prompt for a given example. +# +# Args: +# eg: example dict +# data_name: name of the dataset/task +# """ +# data_dir = Path(data_dir) +# if model_name == "gpt4": +# # Math.Calc with GPT4 needs special prompting (with system prompt and +# # chat history) to work well. +# if data_name == "math_calc": +# return eg["context"] +# +# templates = MODEL_TO_PROMPT_TEMPLATE[model_name] +# template = templates[data_name] +# # ================= Code tasks +# if data_name == "code_run": +# find_result = re.findall(r"func_[0-9]+\(\-?[0-9]+\)", eg['input']) +# func_call = find_result[0] +# func = func_call.split("(")[0] +# return template.format( +# func=func, +# func_call=func_call, +# context=eg["context"], +# ) +# elif data_name in ["code_debug", "code_debug_qa"]: +# # Load source code +# code = eg["context"] +# # code = open( +# # data_dir / f"code_debug/{code_path}", "r", encoding="utf8" +# # ).read() +# if data_name == "code_debug": +# return template.format( +# context=code, +# OPTION_A=eg["options"][0], +# OPTION_B=eg["options"][1], +# OPTION_C=eg["options"][2], +# OPTION_D=eg["options"][3], +# ) +# return template.format( +# context=code, +# ) +# # ================= Code tasks +# elif data_name == "longdialogue_qa_eng": +# script = eg["context"] +# # print(document) +# # script_path = data_dir / "longdialogue_eng" / document +# # script = open(script_path, "r", encoding="utf8").read() +# prompt = template.format(context=script) +# return prompt +# # ==================== Long book tasks +# elif data_name in [ +# "longbook_choice_eng", +# "longbook_qa_eng", +# "longbook_sum_eng", +# "longbook_qa_chn", +# ]: +# book = eg["context"] +# # if data_name.endswith("_eng"): +# # book = open( +# # data_dir / "longbook_eng" / book_path, "r", encoding="utf8" +# # ).read() +# # elif data_name.endswith("_chn"): +# # book = open( +# # data_dir / "longbook_chn" / book_path, "r", encoding="utf8" +# # ).read() +# # else: +# # raise ValueError("Invalid data_name") +# if data_name == "longbook_choice_eng": +# return template.format( +# question=eg["input"], +# context=book, +# OPTION_A=eg["options"][0], +# OPTION_B=eg["options"][1], +# OPTION_C=eg["options"][2], +# OPTION_D=eg["options"][3], +# ) +# elif data_name == "longbook_qa_eng": +# return template.format( +# question=eg["input"], +# context=book, +# ) +# elif data_name == "longbook_sum_eng": +# return template.format( +# context=book, +# ) +# elif data_name == "longbook_qa_chn": +# return template.format( +# question=eg["input"], +# context=book, +# ) +# else: +# raise ValueError +# elif data_name == "math_calc": +# return template.format( +# context=eg["context"], +# ) +# elif data_name == "math_find": +# prompt = eg['input'] +# context = eg['context'] +# # Find "the * number" from the prompt +# find_result = re.findall(r"The .+ of", prompt) +# assert find_result, f"Cannot find the target number in {prompt}" +# target_number = find_result[0].lower()[:-3] +# # Replace the number with the answer +# prefix = f"What is {target_number} in the following list?" +# return template.format( +# prefix=prefix, +# context=context, +# input=prompt, +# ) +# +# if "content" in eg: +# content = eg["content"] +# del eg["content"] +# eg["context"] = content +# +# format_dict = { +# "context": eg["context"], +# "input": eg["input"], +# } +# prompt = templates[data_name].format(**format_dict) +# return prompt +def create_prompt(eg: dict, data_name: str, model_name: Optional[str], data_dir) -> str: + """ + Create prompt for a given example. + + Args: + eg: example dict + data_name: name of the dataset/task + model_name: optional, used to fetch model-specific templates. + """ + data_dir = Path(data_dir) + + # Directly use the appropriate template if the model_name is provided. + if model_name and model_name in MODEL_TO_PROMPT_TEMPLATE: + templates = MODEL_TO_PROMPT_TEMPLATE[model_name] + template = templates[data_name] + else: + # If no model-specific template, return a basic prompt or handle differently. + return eg["context"] + + # Now create the prompt based on the template and task data + if data_name == "code_run": + find_result = re.findall(r"func_[0-9]+\(\-?[0-9]+\)", eg['input']) + func_call = find_result[0] + func = func_call.split("(")[0] + return template.format( + func=func, + func_call=func_call, + context=eg["context"], + ) + elif data_name in ["code_debug", "code_debug_qa"]: + code = eg["context"] + if data_name == "code_debug": + return template.format( + context=code, + OPTION_A=eg["options"][0], + OPTION_B=eg["options"][1], + OPTION_C=eg["options"][2], + OPTION_D=eg["options"][3], + ) + return template.format(context=code) + elif data_name == "longdialogue_qa_eng": + script = eg["context"] + prompt = template.format(context=script) + return prompt + elif data_name in [ + "longbook_choice_eng", + "longbook_qa_eng", + "longbook_sum_eng", + "longbook_qa_chn", + ]: + book = eg["context"] + if data_name == "longbook_choice_eng": + return template.format( + question=eg["input"], + context=book, + OPTION_A=eg["options"][0], + OPTION_B=eg["options"][1], + OPTION_C=eg["options"][2], + OPTION_D=eg["options"][3], + ) + elif data_name == "longbook_qa_eng": + return template.format( + question=eg["input"], + context=book, + ) + elif data_name == "longbook_sum_eng": + return template.format(context=book) + elif data_name == "longbook_qa_chn": + return template.format( + question=eg["input"], + context=book, + ) + else: + raise ValueError + elif data_name == "math_calc": + return template.format(context=eg["context"]) + elif data_name == "math_find": + prompt = eg['input'] + context = eg['context'] + find_result = re.findall(r"The .+ of", prompt) + assert find_result, f"Cannot find the target number in {prompt}" + target_number = find_result[0].lower()[:-3] + prefix = f"What is {target_number} in the following list?" + return template.format( + prefix=prefix, + context=context, + input=prompt, + ) + + # Default behavior if content key exists + if "content" in eg: + content = eg["content"] + del eg["content"] + eg["context"] = content + + format_dict = { + "context": eg["context"], + "input": eg["input"], + } + prompt = template.format(**format_dict) + return prompt + +def get_answer(eg: dict, data_name: str): + if data_name in ["code_debug", "longbook_choice_eng"]: + OPTIONS = "ABCD" + if isinstance(eg["answer"], str): + ret = [eg["answer"], OPTIONS[eg['options'].index(eg["answer"])]] + elif isinstance(eg["answer"], list): + if len(eg["answer"]) == 1: + ret = [eg["answer"][0], OPTIONS[eg['options'].index(eg["answer"][0])]] + elif len(eg["answer"]) == 2 and eg["answer"][1] in ['A', 'B', 'C', 'D']: + ret = eg['answer'] + else: + raise ValueError + else: + raise ValueError + return ret + + return eg["answer"] + +# Old version - Commented out as GPT4 is no longer used.... +# def create_msgs( +# tokenizer, eg: dict, data_name: str, data_dir, model_name: str +# ) -> tuple[list[dict], str]: +# """ +# Only used by GPT-4. +# """ +# prompt = create_prompt(eg, data_name, model_name, data_dir) +# tokens = tokenizer.encode(prompt) +# # - 1000 to have space for system message and other stuff. +# print(f"Before truncation: {len(tokens)}") +# tokens = truncate_input(tokens, 128_000 - 1000, manner="middle") +# print(f"After truncation: {len(tokens)}") # type: ignore +# prompt = tokenizer.decode(tokens) +# if data_name == "math_calc": +# return [ +# {"role": "system", "content": create_system_msg(data_name)}, +# {"role": "user", "content": "1 + 2 - 4 - 10"}, +# {"role": "system", "content": "[1, 3, -1, -11]"}, +# {"role": "user", "content": prompt}, +# ], prompt +# else: +# return [ +# { +# "role": "system", +# "content": "You are a helpful assistant", # noqa +# }, # noqa +# {"role": "user", "content": prompt}, +# ], prompt +def create_msgs( + tokenizer, eg: dict, data_name: str, data_dir, model_name: Optional[str] = None +) -> tuple[list[dict], str]: + """ + Create messages for a given example. + """ + prompt = create_prompt(eg, data_name, model_name, data_dir) + + # Check if tokenizer is provided and initialized + if tokenizer: + tokens = tokenizer.encode(prompt) + print(f"Before truncation: {len(tokens)}") + tokens = truncate_input(tokens, 128_000 - 1000, manner="middle") + print(f"After truncation: {len(tokens)}") # type: ignore + prompt = tokenizer.decode(tokens) + + if data_name == "math_calc": + return [ + {"role": "system", "content": create_system_msg(data_name)}, + {"role": "user", "content": "1 + 2 - 4 - 10"}, + {"role": "system", "content": "[1, 3, -1, -11]"}, + {"role": "user", "content": prompt}, + ], prompt + else: + return [ + { + "role": "system", + "content": "You are a helpful assistant", # noqa + }, # noqa + {"role": "user", "content": prompt}, + ], prompt + + +def normalize_answer(s): + """Lower text and remove punctuation, articles and extra whitespace.""" + + def remove_articles(text): + return re.sub(r"\b(a|an|the)\b", " ", text) + + def white_space_fix(text): + return " ".join(text.split()) + + def remove_punc(text): + exclude = set(string.punctuation) + return "".join(ch for ch in text if ch not in exclude) + + def lower(text): + return text.lower() + + return white_space_fix(remove_articles(remove_punc(lower(s)))) + + +def normalize_zh_answer(s): + """Lower text and remove punctuation, extra whitespace.""" + + def white_space_fix(text): + return "".join(text.split()) + + def remove_punc(text): + cn_punctuation = "！？｡。＂＃＄％＆＇（）＊＋，－／：；＜＝＞＠［＼］＾＿｀｛｜｝～｟｠｢｣､、〃》「」『』【】〔〕〖〗〘〙〚〛〜〝〞〟〰〾〿–—‘’‛“”„‟…‧﹏." # noqa + all_punctuation = set(string.punctuation + cn_punctuation) + return "".join(ch for ch in text if ch not in all_punctuation) + + def lower(text): + return text.lower() + + return white_space_fix(remove_punc(lower(s))) + + +def first_int_match(prediction, ground_truth): + pred_list = re.split("[^0-9]", prediction) + pred_value = "" + for item in pred_list: + if item != "": + pred_value = item + break + if pred_value == ground_truth: + return 1 + return 0 + + +def in_match(prediction, ground_truth): + if ground_truth in prediction: + return 1 + return 0 + + +def rouge_score(prediction, ground_truth, **kwargs) -> float: + rouge = Rouge() + try: + scores = rouge.get_scores([prediction], [ground_truth], avg=True) + except: # noqa + return 0.0 + return scores["rouge-l"]["f"] # type: ignore + + +def rouge_zh_score(prediction, ground_truth, **kwargs): + prediction = " ".join(list(jieba.cut(prediction, cut_all=False))) + ground_truth = " ".join(list(jieba.cut(ground_truth, cut_all=False))) + score = rouge_score(prediction, ground_truth) + return score + + +def f1_score(prediction, ground_truth, **kwargs): + common = Counter(prediction) & Counter(ground_truth) + num_same = sum(common.values()) + if num_same == 0: + return 0 + precision = 1.0 * num_same / len(prediction) + recall = 1.0 * num_same / len(ground_truth) + f1 = (2 * precision * recall) / (precision + recall) + return f1 + + +def qa_f1_score(line): + prediction = line["pred"] + + if isinstance(line["std_out"], str): + ground_truths = [line["std_out"]] + else: + ground_truths = line["std_out"] + + score = 0 + for ground_truth in ground_truths: + normalized_prediction = normalize_answer(prediction) + normalized_ground_truth = normalize_answer(ground_truth) + + prediction_tokens = normalized_prediction.split() + ground_truth_tokens = normalized_ground_truth.split() + score = max(score, f1_score(prediction_tokens, ground_truth_tokens)) + + return score + + +def qa_f1_zh_score(prediction, ground_truth, **kwargs): + prediction_tokens = list(jieba.cut(prediction, cut_all=False)) + ground_truth_tokens = list(jieba.cut(ground_truth, cut_all=False)) + prediction_tokens = [ + normalize_zh_answer(token) for token in prediction_tokens + ] + ground_truth_tokens = [ + normalize_zh_answer(token) for token in ground_truth_tokens + ] + prediction_tokens = [ + token for token in prediction_tokens if len(token) > 0 + ] + ground_truth_tokens = [ + token for token in ground_truth_tokens if len(token) > 0 + ] + return f1_score(prediction_tokens, ground_truth_tokens) + + +def truncate_input(input, max_length, manner="middle"): + if len(input) <= max_length: + return input + if manner == "middle": + return input[0 : max_length // 2] + input[-max_length // 2 :] + else: + return None + + +def load_comprehensive_config(): + # Get the directory of the current script + current_dir = os.path.dirname(os.path.abspath(__file__)) + # Construct the path to the config file + config_path = os.path.join(current_dir, 'Config_Files', 'config.txt') + # Read the config file + config = configparser.ConfigParser() + # Read the configuration file + files_read = config.read(config_path) + if not files_read: + raise FileNotFoundError(f"Config file not found at {config_path}") + return config + + +# FIXME - update to include prompt path in return statement +def load_and_log_configs(): + try: + config = load_comprehensive_config() + if config is None: + logging.error("Config is None, cannot proceed") + return None + # API Keys + anthropic_api_key = config.get('API', 'anthropic_api_key', fallback=None) + logging.debug( + f"Loaded Anthropic API Key: {anthropic_api_key[:5]}...{anthropic_api_key[-5:] if anthropic_api_key else None}") + + cohere_api_key = config.get('API', 'cohere_api_key', fallback=None) + logging.debug( + f"Loaded Cohere API Key: {cohere_api_key[:5]}...{cohere_api_key[-5:] if cohere_api_key else None}") + + groq_api_key = config.get('API', 'groq_api_key', fallback=None) + logging.debug(f"Loaded Groq API Key: {groq_api_key[:5]}...{groq_api_key[-5:] if groq_api_key else None}") + + openai_api_key = config.get('API', 'openai_api_key', fallback=None) + logging.debug( + f"Loaded OpenAI API Key: {openai_api_key[:5]}...{openai_api_key[-5:] if openai_api_key else None}") + + huggingface_api_key = config.get('API', 'huggingface_api_key', fallback=None) + logging.debug( + f"Loaded HuggingFace API Key: {huggingface_api_key[:5]}...{huggingface_api_key[-5:] if huggingface_api_key else None}") + + openrouter_api_key = config.get('API', 'openrouter_api_key', fallback=None) + logging.debug( + f"Loaded OpenRouter API Key: {openrouter_api_key[:5]}...{openrouter_api_key[-5:] if openrouter_api_key else None}") + + deepseek_api_key = config.get('API', 'deepseek_api_key', fallback=None) + logging.debug( + f"Loaded DeepSeek API Key: {deepseek_api_key[:5]}...{deepseek_api_key[-5:] if deepseek_api_key else None}") + + mistral_api_key = config.get('API', 'mistral_api_key', fallback=None) + logging.debug( + f"Loaded Mistral API Key: {mistral_api_key[:5]}...{mistral_api_key[-5:] if mistral_api_key else None}") + + # Models + anthropic_model = config.get('API', 'anthropic_model', fallback='claude-3-sonnet-20240229') + cohere_model = config.get('API', 'cohere_model', fallback='command-r-plus') + groq_model = config.get('API', 'groq_model', fallback='llama3-70b-8192') + openai_model = config.get('API', 'openai_model', fallback='gpt-4-turbo') + huggingface_model = config.get('API', 'huggingface_model', fallback='CohereForAI/c4ai-command-r-plus') + openrouter_model = config.get('API', 'openrouter_model', fallback='microsoft/wizardlm-2-8x22b') + deepseek_model = config.get('API', 'deepseek_model', fallback='deepseek-chat') + mistral_model = config.get('API', 'mistral_model', fallback='mistral-large-latest') + + logging.debug(f"Loaded Anthropic Model: {anthropic_model}") + logging.debug(f"Loaded Cohere Model: {cohere_model}") + logging.debug(f"Loaded Groq Model: {groq_model}") + logging.debug(f"Loaded OpenAI Model: {openai_model}") + logging.debug(f"Loaded HuggingFace Model: {huggingface_model}") + logging.debug(f"Loaded OpenRouter Model: {openrouter_model}") + logging.debug(f"Loaded Deepseek Model: {deepseek_model}") + logging.debug(f"Loaded Mistral Model: {mistral_model}") + + # Local-Models + kobold_api_ip = config.get('Local-API', 'kobold_api_IP', fallback='http://127.0.0.1:5000/api/v1/generate') + kobold_api_key = config.get('Local-API', 'kobold_api_key', fallback='') + + llama_api_IP = config.get('Local-API', 'llama_api_IP', fallback='http://127.0.0.1:8080/v1/chat/completions') + llama_api_key = config.get('Local-API', 'llama_api_key', fallback='') + + ooba_api_IP = config.get('Local-API', 'ooba_api_IP', fallback='http://127.0.0.1:5000/v1/chat/completions') + ooba_api_key = config.get('Local-API', 'ooba_api_key', fallback='') + + tabby_api_IP = config.get('Local-API', 'tabby_api_IP', fallback='http://127.0.0.1:5000/api/v1/generate') + tabby_api_key = config.get('Local-API', 'tabby_api_key', fallback=None) + tabby_model = config.get('services', 'tabby_model', fallback=None) + + vllm_api_url = config.get('Local-API', 'vllm_api_IP', fallback='http://127.0.0.1:500/api/v1/chat/completions') + vllm_api_key = config.get('Local-API', 'vllm_api_key', fallback=None) + vllm_model = config.get('Local-API', 'vllm_model', fallback=None) + + ollama_api_url = config.get('Local-API', 'ollama_api_IP', fallback='http://127.0.0.1:11434/api/generate') + ollama_api_key = config.get('Local-API', 'ollama_api_key', fallback=None) + ollama_model = config.get('Local-API', 'ollama_model', fallback=None) + + aphrodite_api_url = config.get('Local-API', 'aphrodite_api_IP', fallback='http://127.0.0.1:8080/v1/chat/completions') + aphrodite_api_key = config.get('Local-API', 'aphrodite_api_key', fallback='') + + logging.debug(f"Loaded Kobold API IP: {kobold_api_ip}") + logging.debug(f"Loaded Llama API IP: {llama_api_IP}") + logging.debug(f"Loaded Ooba API IP: {ooba_api_IP}") + logging.debug(f"Loaded Tabby API IP: {tabby_api_IP}") + logging.debug(f"Loaded VLLM API URL: {vllm_api_url}") + + # Retrieve output paths from the configuration file + output_path = config.get('Paths', 'output_path', fallback='results') + logging.debug(f"Output path set to: {output_path}") + + # Retrieve processing choice from the configuration file + processing_choice = config.get('Processing', 'processing_choice', fallback='cpu') + logging.debug(f"Processing choice set to: {processing_choice}") + + # Prompts - FIXME + prompt_path = config.get('Prompts', 'prompt_path', fallback='prompts.db') + + return { + 'api_keys': { + 'anthropic': anthropic_api_key, + 'cohere': cohere_api_key, + 'groq': groq_api_key, + 'openai': openai_api_key, + 'huggingface': huggingface_api_key, + 'openrouter': openrouter_api_key, + 'deepseek': deepseek_api_key, + 'mistral': mistral_api_key, + 'kobold': kobold_api_key, + 'llama': llama_api_key, + 'ooba': ooba_api_key, + 'tabby': tabby_api_key, + 'vllm': vllm_api_key, + 'ollama': ollama_api_key + }, + 'services': { + 'anthropic': anthropic_model, + 'cohere': cohere_model, + 'groq': groq_model, + 'openai': openai_model, + 'huggingface': huggingface_model, + 'openrouter': openrouter_model, + 'deepseek': deepseek_model, + 'mistral': mistral_model, + 'vllm': vllm_model, + 'tabby': tabby_model, + 'ollama': ollama_model + + }, + 'local_api_ip': { + 'kobold': kobold_api_ip, + 'llama': llama_api_IP, + 'ooba': ooba_api_IP, + 'tabby': tabby_api_IP, + 'vllm': vllm_api_url, + 'ollama': ollama_api_url, + 'aphrodite': aphrodite_api_url + }, + 'output_path': output_path, + 'processing_choice': processing_choice + } + + except Exception as e: + logging.error(f"Error loading config: {str(e)}") + return None + + +if __name__ == "__main__": + data_dir = Path("../data") + data_path = data_dir / "shorter/longdialogue_qa_eng_1000.jsonl" + examples = list(iter_jsonl(data_path)) + prompt = create_prompt(examples[10], 'longdialogue_qa_eng', 'kimi', data_dir) + print(prompt) diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/prompt.py b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..ba026e86b00a3f8ffc54d8c3d7bc193d6a778062 --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/prompt.py @@ -0,0 +1,62 @@ +gpt4_templates = { + "passkey": "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.\n\n{context}\n\n{input}", # noqa + "number_string": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n\n{input}", # noqa + "kv_retrieval": "Extract the value corresponding to the specified key in the JSON object below.\n\n{context}\n\n{input}", # noqa + # "longbook_sum_eng": "Summarize the book below:\n\n{context}", # noqa + "longbook_qa_eng": "Read the book below and answer a question.\n\n{context}\n\nQuestion: {question}\n\nBe very concise.", # noqa + "longbook_choice_eng": "Read the book and answer the question.\n\n{context}\n\nQuestion: {question}\n\nOnly one of the following options is correct, tell me the answer using one single letter (A, B, C, or D). Don't say anything else.\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}", # noqa + "longbook_sum_eng": "Summarize the following book.\n\n{context}", # noqa + "longbook_qa_chn": "请根据以下书籍回答我的问题。\n\n{context}\n\n问题：{question}\n请尽量简短地回答。", # noqa + "math_find": "{prefix}\n\n{context}\n\n{input}", + "math_calc": "Compute the intermediate values in the following long expression.\n\n{context}", # noqa + "code_run": "Following is a set of Python functions. There is a function called named {func}.\n\n{context}\n\nPlease give me the exact number of the return value of {func_call}. Be concise. Your response must end with the final returned value.", # noqa + "code_debug": "There is ONLY ONE function in the large project that is deliberately made to include an obvious error. Please find the function that contains the most obvious errors. I will give you four options to narrow your scope. You can inspect the options and think. Eventually, tell me the answer using one single letter (A, B, C, or D).\n\n{context}\n\nWhich funtion has deliberate error?\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}\n\nYou should first find the functions in the options. Repeat their content, inspect through code, and at last give me your answer for the function that has the deliberate and obvious error in A, B, C, or D.", # noqa + "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is.\n\nThe dialogue:\n\n---\n\n{context}\n\n---\n\nEnd of dialogue.\n\nWhich character is most likely \"$$MASK$$\"? Just say the name used by the scriptwriter (before the colon marks) of one single character and nothing else.", # noqa +} + +yarn_mistral_templates = { + "passkey": "There is an important info hidden inside a lot of irrelevant text. Find it and memorize it. I will quiz you about the important information.\n\n{context}\n\n{input}\n\nThe pass key is", # noqa + "number_string": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n\n{input}\n\nThe sequence of digits is", # noqa + "kv_retrieval": "Extract the value corresponding to the specified key in the JSON object below.\n\n{context}\n\n{input}", # noqa + "longbook_sum_eng": "Summarize the book below.\n\n{context}\n\nSummary:", # noqa + "longbook_choice_eng": "Read the book and answer the question.\n\n{context}\n\nQuestion: {question}\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}\n\nThe letter of the correct answer is", # noqa + "longbook_qa_eng": "Read the book and answer the question. Be very concise in your answer.\n\n{context}\n\nQuestion: {question}\nAnswer:", # noqa + "longbook_qa_chn": "阅读以下书籍然后回答问题。\n\n{context}\n\n问题：{question}\n答案：", # noqa + "math_find": "{prefix}\n\n{context}\n\n{input}", + "math_calc": "Let us calculate the intermediate values of an expression.\n\nExpression: 1 + 3 + 4\nValues: [1, 4, 8]\n\nExpression: 8 - 3 + 2 - 4\nValues: [8, 5, 7, 3]\n\nExpression: {context}\nValues:", # noqa + "code_run": "There is a function called {func} in the following Python code.\n\n{context}\n\nPlease compute the exact value of {func_call}. The value of {func_call} is", # noqa + "code_debug": "Following is a Python code where exactly one of the functions/methods has a deliberate error that makes it crash.\n\n{context}\n\nOptions:\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}\n\nThe correct option is:", # noqa + "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is.\n\n{context}\n\nThe name that has been replaced with $$MASK$$ is likely", # noqa +} + +claude2_templates = { + "passkey": "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.\n\n{context}\n{input}\nThe pass key is", + "number_string": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n{input}\nThe sequence of digits is", # noqa + "kv_retrieval": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n{input}", + "longbook_sum_eng": "Summarize the following book.\n\n{context}", # noqa + "longbook_choice_eng": "Read the book and answer the question.\n\n{context}\n\nQuestion: {question}\n\nOnly one of the following options is correct, tell me the answer using one single letter (A, B, C, or D). Don't say anything else.\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}", # noqa + "longbook_qa_eng": "Read the novel below and answer a question:\n\n{context}\n\n{input}\nPlease answer as short as possible. The answer is: ", # noqa + "longbook_qa_chn": "请根据以下书籍回答我的问题。\n\n{context}\n\n问题：{question}\n请尽量简短地回答。", # noqa + "math_find": "{prefix}\n\n{context}\n\n{input}", + "math_calc": "Let us calculate the intermediate values of an expression.\nExpression: 1 + 3 + 4\nValues: [1, 4, 8]\n\nExpression: 8 - 3 + 2 - 4\nValues: [8, 5, 7, 3]\n\nExpression: {context}\nValues:", # noqa + "code_run": "In the file functions_module.py, there is a function called ${func}.\n\n\nHere is the content of functions_module.py:\n{context}\n\nPlease give me the exact number of the return value of {func_call}. Your response should end with the sentence \'The return value is:\'.", # noqa + "code_debug": "There is ONLY ONE function in the large project that is deliberately made to include an obvious error. Please find the function that contains the most obvious errors. I will give you four options to narrow your scope. You can inspect through the options and think. Eventually, tell me the answer using one single letter (A, B, C, or D).\n\n{context}\n\nWhich funtion has deliberate error?\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}\n\nYou should first find the functions in the options. Repeat their content, inspect through code, and at last give me your answer for the function that has the deliberate and obvious error in A, B, C, or D.", # noqa + "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is.\n\nThe dialogue:\n\n---\n\n{context}\n\n---\n\nEnd of dialogue.\n\nWhich character is most likely \"$$MASK$$\"? Just say the name used by the scriptwriter (before the colon marks) of one single character and nothing else.", # noqa +} + +kimi_templates = { + "passkey": "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.\n\n{context}\n{input}\nThe pass key is", # noqa + "number_string": "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n\n{context}\n{input}\nThe sequence of digits is", # noqa + "kv_retrieval": "Extract the value corresponding to the specified key in the JSON object below.\n\n{context}\n{input}", # noqa + "longbook_sum_eng": "Summarize the book below:\n\n{file:{context}}", # noqa + "longbook_choice_eng": "Read the book and answer the question.\n\nQuestion: {question}\n\nOnly one of the following options is correct, tell me the answer using one single letter (A, B, C, or D). Don't say anything else.\nA. {OPTION_A}\nB. {OPTION_B}\nC. {OPTION_C}\nD. {OPTION_D}" + "{file:{document}}", # noqa + "longbook_qa_eng": "Read the book below and answer a question.\n\nQuestion: {question}\n\nBe very concise." + "{file:{context}}", # noqa + "longbook_qa_chn": "阅读以下书籍然后回答问题。\n\n问题：{question}\n答案：" + "{file:{context}}", # noqa + "math_find": "{prefix}\n\n{context}\n\n{input}", + "math_calc": "Let us calculate the intermediate values of an expression.\nExpression: 1 + 3 + 4\nValues: [1, 4, 8]\n\nExpression: 8 - 3 + 2 - 4\nValues: [8, 5, 7, 3]\n\nExpression: {context}\nValues:", # noqa + "code_run": "In the file functions_module.py, there is a function called ${func}.\n\n\nHere is the content of functions_module.py:\n\nPlease give me the exact number of the return value of ${func_call}. Your response should end with the sentence 'The return value is:'." + "{context}", # noqa + "code_debug": "Below is a code repository where there is one single function with bugs that causes an error. Please tell me the name of that function.\nWhich function has bugs? Give me the final answer in this format: \"[FINAL ANSWER: XXX]\". Don't say anything else." + "{fcontext}", # noqa + # "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is.\n\nThe name that has been replaced with $$MASK$$ is likely" + "{context}", # noqa + "longdialogue_qa_eng": "Below is a dialogue script where one random occurrence of a character name is replaced with \"$$MASK$$\", and you should try to guess who that character is. Give me the answer using the name before the colons, don't say anything else.\n\n{context}", # noqa +} + diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/test_chat_API_Calls.py b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/test_chat_API_Calls.py new file mode 100644 index 0000000000000000000000000000000000000000..b839087777e83d92cdc1aea53e49866c780ff4ba --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/InifiniteBench/test_chat_API_Calls.py @@ -0,0 +1,106 @@ +# test_chat_API_Calls.py +# Test file for testing the integration of the LLM API calls with the Chat APIs. +# +# Usage: +# python -m unittest test_chat_API_Calls.py + +import unittest + +from LLM_API_Calls import ( + chat_with_openai, + chat_with_anthropic, + chat_with_cohere, + chat_with_groq, + chat_with_openrouter, + chat_with_huggingface, + chat_with_deepseek, + chat_with_mistral +) +from eval_utils import load_and_log_configs + + +class TestLLMAPICallsIntegration(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.config = load_and_log_configs() + if cls.config is None: + raise ValueError("Failed to load configuration") + + def test_chat_with_openai(self): + api_key = self.config['api_keys'].get('openai') + model = self.config['services'].get('openai') + if not api_key: + self.skipTest("OpenAI API key not available") + response = chat_with_openai(api_key, "Hello, how are you?", "Respond briefly", temp=0.7, system_message="You are a helpful assistant.") + print("OpenAI Response: " + response + "\n") + self.assertIsInstance(response, str) + self.assertTrue(len(response) > 0) + + def test_chat_with_anthropic(self): + api_key = self.config['api_keys'].get('anthropic') + model = self.config['services'].get('anthropic') + if not api_key: + self.skipTest("Anthropic API key not available") + response = chat_with_anthropic(api_key, "Hello, how are you?", model, "Respond briefly") + print("Anthropic Response: " + response + "\n") + self.assertIsInstance(response, str) + self.assertTrue(len(response) > 0) + + def test_chat_with_cohere(self): + api_key = self.config['api_keys'].get('cohere') + model = self.config['services'].get('cohere') + if not api_key: + self.skipTest("Cohere API key not available") + response = chat_with_cohere(api_key, "Hello, how are you?", model, "Respond briefly") + print("Cohere Response: " + response + "\n") + self.assertIsInstance(response, str) + self.assertTrue(len(response) > 0) + + def test_chat_with_groq(self): + api_key = self.config['api_keys'].get('groq') + if not api_key: + self.skipTest("Groq API key not available") + response = chat_with_groq(api_key, "Hello, how are you?", "Respond briefly") + print("Groq Response: " + response + "\n") + self.assertIsInstance(response, str) + self.assertTrue(len(response) > 0) + + def test_chat_with_openrouter(self): + api_key = self.config['api_keys'].get('openrouter') + if not api_key: + self.skipTest("OpenRouter API key not available") + response = chat_with_openrouter(api_key, "Hello, how are you?", "Respond briefly") + print("OpenRouter Response: " + response + "\n") + self.assertIsInstance(response, str) + self.assertTrue(len(response) > 0) + + def test_chat_with_huggingface(self): + api_key = self.config['api_keys'].get('huggingface') + if not api_key: + self.skipTest("HuggingFace API key not available") + response = chat_with_huggingface(api_key, "Hello, how are you?", "Respond briefly") + print("Huggingface Response: " + response + "\n") + self.assertIsInstance(response, str) + self.assertTrue(len(response) > 0) + + def test_chat_with_deepseek(self): + api_key = self.config['api_keys'].get('deepseek') + if not api_key: + self.skipTest("DeepSeek API key not available") + response = chat_with_deepseek(api_key, "Hello, how are you?", "Respond briefly") + print("DeepSeek Response: " + response + "\n") + self.assertIsInstance(response, str) + self.assertTrue(len(response) > 0) + + def test_chat_with_mistral(self): + api_key = self.config['api_keys'].get('mistral') + if not api_key: + self.skipTest("Mistral API key not available") + response = chat_with_mistral(api_key, "Hello, how are you?", "Respond briefly") + print("Mistral Response: " + response + "\n") + self.assertIsInstance(response, str) + self.assertTrue(len(response) > 0) + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/README.md b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/README.md new file mode 100644 index 0000000000000000000000000000000000000000..7417c6361677819fe210e94de3568201db180d9d --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/README.md @@ -0,0 +1,200 @@ +
+ +
+
+ +# InfiniteBench: Extending Long Context Evaluation Beyond 100K Tokens + +

+ 中文 • + English • + Paper +

+ +
+ +## Introduction + +Welcome to InfiniteBench, a cutting-edge benchmark tailored for evaluating the capabilities of language models to process, understand, and reason over super long contexts (100k+ tokens). Long contexts are crucial for enhancing applications with LLMs and achieving high-level interaction. InfiniteBench is designed to push the boundaries of language models by testing them against a context length of 100k+, which is 10 times longer than traditional datasets. + +## Features + +- **Loooong Context:** InfiniteBench is a pioneer in testing language models with a context length of 100k+, offering an unparalleled challenge in the field. +- **Diverse Domain:** The benchmark comprises 12 unique tasks, each crafted to assess different aspects of language processing and comprehension in extended contexts. +- **Specialized Test:** InfiniteBench consists of tasks that state-of-the-art LLMs are known to be capable of when using shorter context. This ensures that the performance degradation is only caused by the length of the contexts. +- **Real-World and Synthetic Scenarios:** The tasks are a mix of real-world scenarios and synthetic constructs, ensuring a comprehensive evaluation of models. Real-world scenarios make the test pragmatic, and synthetic ones leave the space for extending the context length further with ease. + +## Task Composition + +
+ +
+ +| Task Name | Context | # Examples | Avg Input Tokens | Avg Output Tokens | Description | +| -------------------- | ------------- | ---------- | ---------------- | ----------------- | ------------------------------------------------------------------------------------------- | +| En.Sum | Fake Book | 103 | 171.5k | 1.1k | Summarization of a fake book created with core entity substitution. | +| En.QA | Fake Book | 351 | 192.6k | 4.8 | Free-form question answering based on the fake book. | +| En.MC | Fake Book | 229 | 184.4k | 5.3 | Multiple choice questions derived from the fake book. | +| En.Dia | Script | 200 | 103.6k | 3.4 | Identification of talkers in partially anonymized scripts. | +| Zh.QA | New Book | 175 | 2068.6k | 6.3 | Question answering on a set of newly collected books. | +| Code.Debug | Code Document | 394 | 114.7k | 4.8 | Finding which function in a code repo contains an crashing error (in multiple choice form). | +| Code.Run | Synthetic | 400 | 75.2k | 1.3 | Simulating execution of multiple simple, synthetic functions. | +| Math.Calc | Synthetic | 50 | 43.9k | 43.9k | Calculations involving super-long arithmetic equations. | +| Math.Find | Synthetic | 350 | 87.9k | 1.3 | Finding special integers in a lengthy list. | +| Retrieve.PassKey[^1] | Synthetic | 590 | 122.4k | 2.0 | Retrieving hidden keys in a noisy long context. | +| Retrieve.Number | Synthetic | 590 | 122.4k | 4.0 | Locating repeated hidden numbers in a noisy long context. | +| Retrieve.KV[^2] | Synthetic | 500 | 89.9k | 22.7 | Finding the corresponding value from a dictionary and a key. | + +## How to Download Data + +Click here to download data from 🤗 Huggingface directly: + +### Using 🤗 Datasets + +Alternatively, you can download using the 🤗 Datasets library as follows. + +```python +from datasets import load_dataset, Value, Sequence +ft = Features({"id": Value("int64"), "context": Value("string"), "input": Value("string"), "answer": Sequence(Value("string")), "options": Sequence(Value("string"))}) +dataset = load_dataset("xinrongzhang2022/InfiniteBench", features=ft) +``` +### Using Scripts + +```shell +cd InfiniteBench +bash scripts/download_dataset.sh +``` + +This will directly dump the data to `data`. + +## Evaluation Result + +We evaluate SOTA proprietary and open-source LLMs, the result is as follows. + +| Task Name | GPT-4 | YaRN-Mistral-7B | Kimi-Chat | Claude 2 | Yi-6B-200K | Yi-34B-200K | Chatglm3-6B-128K | +| ---------------- | ------ | --------------- | --------- | -------- | -----------| -----------| -----------| +| Retrieve.PassKey | 100% | 92.71% | 98.14% | 97.80% | 100.00% | 100.00% | 92.20% | +| Retrieve.Number | 100% | 56.61% | 95.42% | 98.14% | 94.92% | 100.00% | 80.68% | +| Retrieve.KV | 89.00% | < 5% | 53.60% | 65.40% | < 5% | < 5% | < 5% | +| En.Sum | 14.73% | 9.09% | 17.96% | 14.50% | < 5% | < 5% |< 5% | +| En.QA | 22.44% | 9.55% | 16.52% | 11.97% | 9.20% | 12.17% |< 5% | +| En.MC | 67.25% | 27.95% | 72.49% | 62.88% | 36.68% |38.43% |10.48% | +| En.Dia | 8.50% | 7.50% | 11.50% | 46.50% | < 5% |< 5% |< 5% | +| Zh.QA | 25.96% | 16.98% | 17.93% | 9.64% | 15.07% |13.61% |< 5% | +| Code.Debug | 37.06% | < 5% | 17.77% | < 5% | 9.14% |13.96% |7.36% | +| Code.Run | 23.25% | < 5% | < 5% | < 5% | < 5% |< 5% |< 5% | +| Math.Calc | < 5% | < 5% | < 5% | < 5% | < 5% |< 5% |< 5% | +| Math.Find | 60.00% | 17.14% | 12.57% | 32.29% | < 5% |25.71% |7.71% | + +Note: + +1. The evaluation code for YaRN-Mistral-7B is implemented by ourselves, and please contact us or submit an issue if there are any problems. +2. Kimi-Chat, Claude 2, and GPT-4 are evaluated using the official API with default configuration. +3. For Math.Calc, the values in the parentheses have a measurement unit of 0.01%. This is because it is easy to get a very low score on this task. +4. The metric for task Math.Find, Math.Calc, Code.Run, Code.Debug, En.Dia, En.MC, Retrieve.KV, Retrieve.Number, and Retrieve.PassKey is accuracy; + + The metric for task Zh.QA and En.QA are ROUGE F1 score; + + The metric for En.Sum is the `rougeLsum` score from the 🤗 Evaluate library. + + + +
+ +
+ +## Installation + +```shell +pip install -r requirements.txt +``` + +## How to Run + +Download the dataset the `data` folder (or set the `--data_dir` argument to the location of the dataset). The data folder structure should be as follows. + +``` +InfiniteBench +├── data +│ ├── code_debug.jsonl +│ ├── code_run.jsonl +│ ├── kv_retrieval.jsonl +│ ├── longbook_choice_eng.jsonl +│ ├── longbook_qa_chn.jsonl +│ ├── longbook_qa_eng.jsonl +│ ├── longbook_sum_eng.jsonl +│ ├── longdialogue_qa_eng.jsonl +│ ├── math_calc.jsonl +│ ├── math_find.jsonl +│ ├── number_string.jsonl +│ ├── passkey.jsonl +│ └── construct_synthetic_dataset.py +... +``` + +Then, in the `src` folder, execute: + +```shell +python eval_yarn_mistral.py --task kv_retrieval +python eval_gpt4.py --task longbook_sum_qa +python eval_rwkv.py --task passkey +``` + +The available tasks are: + +| Task Name | Argument to specify in `--task` | +| ---------------- | ------------------------------- | +| En.Sum | longbook_sum_eng | +| En.QA | longbook_qa_eng | +| En.MC | longbook_choice_eng | +| En.Dia | longdialogue_qa_eng | +| Zh.QA | longbook_qa_chn | +| Code.Debug | code_debug | +| Code.Run | code_run | +| Math.Calc | math_calc | +| Math.Find | math_find | +| Retrieve.PassKey | passkey | +| Retrieve.Number | number_string | +| Retrieve.KV | kv_retrieval | + +## Citation + +> This will be updated when our preprint paper is released. + +```bibtex +@inproceedings{zhang-etal-2024-bench, + title = "$\infty${B}ench: Extending Long Context Evaluation Beyond 100{K} Tokens", + author = "Zhang, Xinrong and + Chen, Yingfa and + Hu, Shengding and + Xu, Zihang and + Chen, Junhao and + Hao, Moo and + Han, Xu and + Thai, Zhen and + Wang, Shuo and + Liu, Zhiyuan and + Sun, Maosong", + editor = "Ku, Lun-Wei and + Martins, Andre and + Srikumar, Vivek", + booktitle = "Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)", + month = aug, + year = "2024", + address = "Bangkok, Thailand", + publisher = "Association for Computational Linguistics", + url = "https://aclanthology.org/2024.acl-long.814", + pages = "15262--15277", + abstract = "Processing and reasoning over long contexts is crucial for many practical applications of Large Language Models (LLMs), such as document comprehension and agent construction. Despite recent strides in making LLMs process contexts with more than 100K tokens, there is currently a lack of a standardized benchmark to evaluate this long-context capability. Existing public benchmarks typically focus on contexts around 10K tokens, limiting the assessment and comparison of LLMs in processing longer contexts. In this paper, we propose , the first LLM benchmark featuring an average data length surpassing 100K tokens. comprises synthetic and realistic tasks spanning diverse domains in English and Chinese. The tasks in are designed to require an understanding of long dependencies in contexts and make simply retrieving a limited number of passages from contexts not sufficient for these tasks. Based on , we evaluate several state-of-the-art LLMs tailored for processing long contexts. The experimental results indicate that existing long-context LLMs still require significant advancements to process 100K+ contexts effectively. Furthermore, we present three intriguing analyses regarding the behavior of LLMs processing long context. Our code and data is released.", +} +``` + +## Acknowledgement + +Thanks to Cong Feng, Zhongwu Zhai, Guoyang Zeng, Chenyang Song, Renjie Luo, Chaoqun He, Yuge Tu, Bowen Ping, Yujie Huang, Yudong Mei, Kaihuo Zhang, Weilin Zhao, Ao Sun, Yulin Chen, Ganqu Cui. + +## References + +[^1]: Mohtashami, Amirkeivan and Martin Jaggi. "Landmark Attention: Random-Access Infinite Context Length for Transformers." ArXiv abs/2305.16300 (2023): n. pag. + +[^2]: Liu, Nelson F. et al. "Lost in the Middle: How Language Models Use Long Contexts." ArXiv abs/2307.03172 (2023): n. pag. diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/README_ZH.md b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/README_ZH.md new file mode 100644 index 0000000000000000000000000000000000000000..23907e32249df3b155b325a7f62fdb3d85590acc --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/README_ZH.md @@ -0,0 +1,172 @@ +
+ +
+
+ +# InfiniteBench: Extending Long Context Evaluation Beyond 100K Tokens + +

+ 中文 • + English • + 论文 +

+ +
+ +## 简介 + +理解、处理长文本，是大模型迈向更深层次理解与交互阶段必备的能力。现已有大模型声称可以处理100k+的长序列，但是对应的标准评测集却是空缺的。为此，我们构建了一个面向 100k+ 的评测集，InfiniteBench。该评测集针对大模型在长文本方面的五项能力而设计：检索、数学、代码、问答、和摘要。 + +## 特点 + +- **长上下文:** InfiniteBench 测试数据的平均上下文长度为195k，远超现有评测数据。 +- **多领域多语言:** InfiniteBench 评测集包含12个任务，包括中英双语，涵盖了检索、数学、代码、问答、和摘要等5个领域。 +- **前瞻性挑战性:** InfiniteBench 测试任务，对标当前最强的模型如 GPT-4, Claude 2 等。 +- **真实场景与合成场景:** InfiniteBench 既包含真实场景数据，探测大模型在处理实际问题的能力；也包含合成数据，为测试数据拓展上下文窗口提供了便捷。 + +## 任务构成 + +| Task Name | Context | # Examples | Avg Input Tokens | Avg Output Tokens | Description | +| -------------------- | ------------- | ---------- | ---------------- | ----------------- | ------------------------------------------------------------------------------------------- | +| En.Sum | Fake Book | 103 | 171.5k | 1.1k | Summarization of a fake book created with core entity substitution. | +| En.QA | Fake Book | 351 | 192.6k | 4.8 | Free-form question answering based on the fake book. | +| En.MC | Fake Book | 229 | 184.4k | 5.3 | Multiple choice questions derived from the fake book. | +| En.Dia | Script | 200 | 103.6k | 3.4 | Identification of talkers in partially anonymized scripts. | +| Zh.QA | New Book | 175 | 2068.6k | 6.3 | Question answering on a set of newly collected books. | +| Code.Debug | Code Document | 394 | 114.7k | 4.8 | Finding which function in a code repo contains an crashing error (in multiple choice form). | +| Code.Run | Synthetic | 400 | 75.2k | 1.3 | Simulating execution of multiple simple, synthetic functions. | +| Math.Calc | Synthetic | 50 | 43.9k | 43.9k | Calculations involving super-long arithmetic equations. | +| Math.Find | Synthetic | 350 | 87.9k | 1.3 | Finding special integers in a lengthy list. | +| Retrieve.PassKey[^1] | Synthetic | 590 | 122.4k | 2.0 | Retrieving hidden keys in a noisy long context. | +| Retrieve.Number | Synthetic | 590 | 122.4k | 4.0 | Locating repeated hidden numbers in a noisy long context. | +| Retrieve.KV[^2] | Synthetic | 500 | 89.9k | 22.7 | Finding the corresponding value from a dictionary and a key. | + + +## 评测结果 + +我们在 SOTA 模型上评测了 InfiniteBench 结果如下： + +| Task Name | GPT-4 | YaRN-Mistral-7B | Kimi-Chat | Claude 2 | Yi-6B-200K | Yi-34B-200K | Chatglm3-6B-128K | +| ---------------- | ------ | --------------- | --------- | -------- | -----------| -----------| -----------| +| Retrieve.PassKey | 100% | 92.71% | 98.14% | 97.80% | 100.00% | 100.00% | 92.20% | +| Retrieve.Number | 100% | 56.61% | 95.42% | 98.14% | 94.92% | 100.00% | 80.68% | +| Retrieve.KV | 89.00% | < 5% | 53.60% | 65.40% | < 5% | < 5% | < 5% | +| En.Sum | 14.73% | 9.09% | 17.96% | 14.50% | < 5% | < 5% |< 5% | +| En.QA | 22.44% | 9.55% | 16.52% | 11.97% | 9.20% | 12.17% |< 5% | +| En.MC | 67.25% | 27.95% | 72.49% | 62.88% | 36.68% |38.43% |10.48% | +| En.Dia | 8.50% | 7.50% | 11.50% | 46.50% | < 5% |< 5% |< 5% | +| Zh.QA | 25.96% | 16.98% | 17.93% | 9.64% | 15.07% |13.61% |< 5% | +| Code.Debug | 37.06% | < 5% | 17.77% | < 5% | 9.14% |13.96% |7.36% | +| Code.Run | 23.25% | < 5% | < 5% | < 5% | < 5% |< 5% |< 5% | +| Math.Calc | < 5% | < 5% | < 5% | < 5% | < 5% |< 5% |< 5% | +| Math.Find | 60.00% | 17.14% | 12.57% | 32.29% | < 5% |25.71% |7.71% | + +注： + +1. YaRN-Mistral-7B 实现代码已开源在仓库，请大家批评指正；Kimi-Chat 和 Claude 2 使用用户界面评测，GPT-4 使用 API 评测，均使用官方默认配置。 + + +## 评测 + +## 获取数据集 + +从 下载数据集到 `infinitebench/data` 路径下（我们将评测数据集放在 InfiniteBench 目录下），得到文件如下： + +``` +InfiniteBench +├── data +│ ├── code_debug.jsonl +│ ├── code_run.jsonl +│ ├── kv_retrieval.jsonl +│ ├── longbook_choice_eng.jsonl +│ ├── longbook_qa_chn.jsonl +│ ├── longbook_qa_eng.jsonl +│ ├── longbook_sum_eng.jsonl +│ ├── longdialogue_qa_eng.jsonl +│ ├── math_calc.jsonl +│ ├── math_find.jsonl +│ ├── number_string.jsonl +│ ├── passkey.jsonl +│ └── construct_synthetic_dataset.py +... +``` + +或者使用 Datasets 下载： + +```python +from datasets import load_dataset, Value, Sequence +ft = Features({"id": Value("int64"), "context": Value("string"), "input": Value("string"), "answer": Sequence(Value("string")), "options": Sequence(Value("string"))}) +dataset = load_dataset("xinrongzhang2022/InfiniteBench", features=ft) +``` + +### 安装依赖 + +```shell +pip install -r requiremnets.txt +``` + +### 推理 + +比如，评测 GPT-4 在 Retrieve.PassKey 任务上的表现： + +```shell +cd src +python eval_gpt4.py --task passkey +``` + +可以选择的 `--task` 有： + +- `passkey` +- `number_string` +- `kv_retrieval` +- `longbook_sum_eng` +- `longbook_qa_eng` +- `longbook_qa_chn` +- `longbook_choice_eng` +- `longdialogue_qa_eng` +- `math_calc` +- `math_find` +- `code_debug` +- `code_run` + +#### 计算分数 + +```shell +python compute_scores.py +``` + +## 引用 + +> This will be updated when our preprint paper is released. + +```bibtex +@inproceedings{zhang-etal-2024-bench, + title = "$\infty${B}ench: Extending Long Context Evaluation Beyond 100{K} Tokens", + author = "Zhang, Xinrong and + Chen, Yingfa and + Hu, Shengding and + Xu, Zihang and + Chen, Junhao and + Hao, Moo and + Han, Xu and + Thai, Zhen and + Wang, Shuo and + Liu, Zhiyuan and + Sun, Maosong", + editor = "Ku, Lun-Wei and + Martins, Andre and + Srikumar, Vivek", + booktitle = "Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)", + month = aug, + year = "2024", + address = "Bangkok, Thailand", + publisher = "Association for Computational Linguistics", + url = "https://aclanthology.org/2024.acl-long.814", + pages = "15262--15277", + abstract = "Processing and reasoning over long contexts is crucial for many practical applications of Large Language Models (LLMs), such as document comprehension and agent construction. Despite recent strides in making LLMs process contexts with more than 100K tokens, there is currently a lack of a standardized benchmark to evaluate this long-context capability. Existing public benchmarks typically focus on contexts around 10K tokens, limiting the assessment and comparison of LLMs in processing longer contexts. In this paper, we propose , the first LLM benchmark featuring an average data length surpassing 100K tokens. comprises synthetic and realistic tasks spanning diverse domains in English and Chinese. The tasks in are designed to require an understanding of long dependencies in contexts and make simply retrieving a limited number of passages from contexts not sufficient for these tasks. Based on , we evaluate several state-of-the-art LLMs tailored for processing long contexts. The experimental results indicate that existing long-context LLMs still require significant advancements to process 100K+ contexts effectively. Furthermore, we present three intriguing analyses regarding the behavior of LLMs processing long context. Our code and data is released.", +} +``` + +## 参考文献 +[^1]: Mohtashami, Amirkeivan and Martin Jaggi. “Landmark Attention: Random-Access Infinite Context Length for Transformers.” ArXiv abs/2305.16300 (2023): n. pag. +[^2]: Liu, Nelson F. et al. “Lost in the Middle: How Language Models Use Long Contexts.” ArXiv abs/2307.03172 (2023): n. pag. diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/__init__.py b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/InfiniteBench/PUT_DATASETS_HERE.txt b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/InfiniteBench/PUT_DATASETS_HERE.txt new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/__init__.py b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/collections.json b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/collections.json new file mode 100644 index 0000000000000000000000000000000000000000..432f6830be65c05e3e13b47dfe27b9642dcb105f --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/collections.json @@ -0,0 +1 @@ +[[843, 181, 649, 974, 531, 402, 1100, 769, 641, 1094, 529, 584, 504, 920, 526, 759, 358, 962, 487, 243, 428, 117, 523, 1032, 924, 814, 739, 754, 804, 683, 949, 901, 732, 256, 824, 861, 494, 972, 996, 280, 130, 768, 469, 457, 945, 940, 317, 985, 268, 18, 334, 327, 370, 166, 207], [21, 278, 89, 633, 559, 516, 851, 830, 637, 626, 958, 123, 813, 249, 698, 757, 976, 556, 896, 802, 73, 1059, 74, 846, 669, 620, 323, 823, 907, 856, 122, 55, 70, 167, 622, 939, 987, 508, 564, 533, 200, 538, 443, 1098, 1029, 627, 731, 829, 330, 444, 960, 692, 363, 1005, 284], [815, 1095, 879, 864, 796, 397, 702, 1093, 677, 114, 1061, 957, 221, 558, 299, 92, 124, 578, 366, 204, 812, 993, 474, 13, 540, 158, 696, 25, 462, 715, 1060, 1089, 596, 997, 116, 657, 863, 58, 413, 819, 825, 353, 269, 873, 125, 880, 422, 934, 19, 827, 890, 886, 678, 505, 340], [319, 310, 1030, 423, 952, 889, 518, 1076, 473, 387, 937, 275, 155, 289, 1091, 590, 287, 30, 770, 244, 361, 594, 906, 176, 1042, 758, 588, 90, 600, 1083, 121, 638, 688, 836, 903, 826, 891, 730, 625, 545, 695, 948, 1013, 706, 747, 69, 718, 860, 364, 205, 1096, 717, 102, 1043, 274], [1000, 308, 492, 845, 98, 915, 910, 820, 242, 301, 699, 493, 429, 272, 565, 382, 1004, 617, 1078, 751, 923, 557, 385, 23, 393, 262, 240, 101, 1090, 36, 1008, 686, 185, 729, 16, 645, 68, 392, 991, 454, 159, 542, 346, 571, 1020, 237, 679, 1049, 303, 685, 8, 1047, 1079, 378, 48], [1077, 32, 521, 367, 15, 432, 1069, 113, 3, 875, 65, 1051, 119, 248, 986, 931, 234, 336, 782, 634, 85, 53, 288, 965, 917, 231, 992, 1099, 644, 723, 838, 463, 1067, 194, 1080, 552, 195, 928, 52, 760, 225, 989, 735, 727, 362, 400, 842, 595, 390, 201, 510, 562, 664, 1053, 88], [1062, 78, 936, 490, 324, 701, 71, 466, 375, 503, 1027, 703, 292, 647, 132, 46, 115, 263, 253, 309, 480, 63, 887, 484, 1054, 911, 514, 871, 662, 658, 693, 134, 456, 821, 963, 28, 351, 550, 118, 335, 441, 543, 832, 348, 153, 892, 847, 857, 978, 661, 943, 675, 245, 541, 955], [188, 403, 137, 5, 705, 549, 611, 94, 650, 401, 561, 208, 405, 233, 302, 872, 983, 297, 445, 673, 828, 228, 927, 357, 199, 532, 1035, 579, 39, 853, 653, 461, 455, 76, 391, 131, 279, 801, 746, 547, 22, 761, 612, 265, 157, 371, 291, 772, 66, 639, 386, 567, 1007, 877, 805], [800, 294, 964, 169, 1031, 618, 979, 1037, 162, 902, 990, 316, 49, 722, 971, 365, 506, 676, 126, 878, 882, 325, 659, 277, 576, 525, 458, 352, 376, 1003, 665, 470, 33, 798, 750, 7, 740, 1010, 572, 1016, 395, 1086, 267, 778, 648, 859, 811, 209, 172, 716, 869, 486, 140, 147, 141], [1021, 286, 670, 721, 973, 707, 495, 154, 1019, 251, 315, 741, 913, 865, 95, 6, 214, 1045, 374, 313, 950, 1044, 198, 953, 99, 840, 789, 672, 527, 406, 866, 787, 681, 276, 954, 14, 674, 12, 599, 912, 694, 610, 434, 555, 320, 548, 792, 369, 756, 143, 1082, 1075, 988, 296, 224]] \ No newline at end of file diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/construct_synthetic_dataset.py b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/construct_synthetic_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..317d16659cebfde3890316ddf9e77d6f247d2cd5 --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/construct_synthetic_dataset.py @@ -0,0 +1,413 @@ +import jsonlines +import random +import os +import re +import importlib.util +import json + + +def build_number_string(): + #####32 + # prompt = "There is an important info hidden inside a lot of irrelevant text. Find it. I will quiz you about the important information there.\n" + #####25 + noise = "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again.\n" + #####26 + ans = "The sequence of digits is {key}. Remember it. {key} is the sequence of digits.\n" + #####10 + question = "What is the sequence of digits?" + + + target_length = [1024 * 64, 1024 * 128] + num_noise = [2610, 5220] + step = [45, 90] + repeat_time = 10 + for i in range(1, 2): + target_length_i = target_length[i] + step_i = step[i] + num_noise_i = num_noise[i] + ret = [] + for j in range(0, num_noise_i+1, step_i): + input_text = noise * j + ans + noise * (num_noise_i - j) + for t in range(repeat_time): + keys = [] + for k in range(5): + keys.append(str(random.randint(0,9))) + for k in range(5): + pos = random.randint(0,5+k-1) + keys.insert(pos, keys[pos]) + key_t = "".join(keys) + ret.append({"context": input_text.replace("{key}", key_t), "answer": key_t, "input": question, "len": 26 * (num_noise_i - j)}) + fw = jsonlines.open("number_string.jsonl", 'w') + fw.write_all(ret) + fw.close() + + +def build_passkey(): + #####32 + # prompt = "There is an important info hidden inside a lot of irrelevant text. Find it and memorize them. I will quiz you about the important information there.\n" + #####25 + noise = "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again.\n" + #####26 + ans = "The pass key is {key}. Remember it. {key} is the pass key.\n" + #####10 + question = "What is the pass key?" + + target_length = [1024 * 8, 1024 * 16, 1024 * 32, 1024 * 64, 1024 * 128, 1024 * 256] + num_noise = [326, 652, 1305, 2610, 5220, 10440] + step = [6,12 ,22, 45, 90, 180] + repeat_time = 5 + for i in range(0,4): + target_length_i = target_length[i] + step_i = step[i] + num_noise_i = num_noise[i] + ret = [] + for j in range(0, num_noise_i+1, step_i): + input_text = noise * j + ans + noise * (num_noise_i - j) + for t in range(repeat_time): + keys = [] + for k in range(5): + keys.append(str(random.randint(0,9))) + + key_t = "".join(keys) + ret.append({"input": question, "context": input_text.replace("{key}", key_t), "answer": key_t, "len": 26 * (num_noise_i - j)}) + fw = jsonlines.open("passkey_%d.jsonl"%target_length_i, 'w') + fw.write_all(ret) + fw.close() + + +def build_kv_retrieval(): + + target_length = [64 * 1024, 128 * 1024] + # interv = [16, 7] + nsample = [500, 500] + nnoise = [928, 2500] + for ii in range(1, 2): + cnt = -1 + ret = [] + + with jsonlines.open("kv-retrieval-3000_keys.jsonl") as fin: + for line in fin: + print(len(line["ordered_kv_records"])) + # return 0 + cnt += 1 + if cnt == nsample[ii]: + break + ans_id = min(int(cnt * nnoise[ii] / nsample[ii]), nnoise[ii]) + + text = "JSON data:\n{" + t = -1 + random.shuffle(line["ordered_kv_records"]) + for item in line["ordered_kv_records"]: + t += 1 + if t == nnoise[ii]: + break + text += "\"" + item[0] + "\": \"" + item[1] + "\", " + text = text[:-2] + '}' + question = "\nKey: \"" + line["ordered_kv_records"][ans_id][0] + "\"\nThe value associated with the specified key is: " + # text += "\nKey: \"" + line["ordered_kv_records"][ans_id][0] + "\"\nThe value associated with the specified key is: " + # print(len(tokenizer.encode(text))) + # break + ret.append({"id": cnt, "context": text, "input": question, "answer": line["ordered_kv_records"][ans_id][1]}) + + + fw = jsonlines.open("kv_retrieval.jsonl", 'w') + fw.write_all(ret) + fw.close() + + +def generate_random_list(length, _min, _max, task): + # random_list = [random.randint(_min, _max) for _ in range(length)] + # ret_list = random_list.copy() + + if task == "largest number": + _max = random.randint(int(_max * 0.8), _max) + random_list = [random.randint(_min, _max) for _ in range(length)] + ret_list = random_list.copy() + ans = max(random_list) + input = str(ret_list) + elif task == "second largest number": + _max = random.randint(int(_max * 0.8), _max) + random_list = [random.randint(_min, _max) for _ in range(length)] + ret_list = random_list.copy() + target = max(random_list) + while target == max(random_list): + random_list.remove(max(random_list)) + ans = max(random_list) + input = str(ret_list) + + elif task == "third largest number": + _max = random.randint(int(_max * 0.8), _max) + random_list = [random.randint(_min, _max) for _ in range(length)] + ret_list = random_list.copy() + target = max(random_list) + while target == max(random_list): + random_list.remove(max(random_list)) + target = max(random_list) + while target == max(random_list): + random_list.remove(max(random_list)) + ans = max(random_list) + input = str(ret_list) + + elif task == "smallest number": + _min = random.randint(_min, int(_max * 0.2)) + random_list = [random.randint(_min, _max) for _ in range(length)] + ret_list = random_list.copy() + ans = min(random_list) + input = str(ret_list) + + elif task == "second smallest number": + _min = random.randint(_min, int(_max * 0.2)) + random_list = [random.randint(_min, _max) for _ in range(length)] + ret_list = random_list.copy() + target = min(random_list) + while target == min(random_list): + random_list.remove(min(random_list)) + ans = min(random_list) + input = str(ret_list) + + elif task == "third smallest number": + _min = random.randint(_min, int(_max * 0.2)) + random_list = [random.randint(_min, _max) for _ in range(length)] + ret_list = random_list.copy() + target = min(random_list) + while target == min(random_list): + random_list.remove(min(random_list)) + target = min(random_list) + while target == min(random_list): + random_list.remove(min(random_list)) + ans = min(random_list) + input = str(ret_list) + elif task == "median": + if random.random() > 0.5: + _min = random.randint(_min, int(_max * 0.2)) + random_list = [random.randint(_min, _max) for _ in range(length)] + else: + _max = random.randint(int(_max * 0.8), _max) + random_list = [random.randint(_min, _max) for _ in range(length)] + ret_list = random_list.copy() + random_list.sort() + if len(random_list)%2 == 1: + ans = random_list[len(random_list)//2] + else: + ans = (random_list[len(random_list)//2] + random_list[len(random_list)//2-1])/2 + input = str(ret_list) + elif task == "expression": + random_list = [random.randint(_min, _max) for _ in range(length)] + ret_list = random_list.copy() + input = str(random_list[0]) + value = random_list[0] + ans = [] + for i in range(1, length): + poss = random.random() + if poss > 0.5: + if value + random_list[i] > _max: + random_list[i] = random.randint(_min, _max-value) + + input += " + " + str(random_list[i]) + value += random_list[i] + + else: + if value - random_list[i] < 0: + random_list[i] = random.randint(_min, value) + input += " - " + str(random_list[i]) + value -= random_list[i] + ans.append(value) + + + else: + print("Invalid task") + ans = None + + return ans, input + + +def generate_math_qa(list_length, min_val, max_val, tasks=None): + num_samples = 50 + ret = [] + prompts = { + "largest number": "Find the largest number from the list below:", + "second largest number": "Find the second largest number from the list below:", + "third largest number": "Find the third largest number from the list below:", + "smallest number": "Find the smallest number from the list below:", + "second smallest number": "Find the second smallest number from the list below:", + "third smallest number": "Find the third smallest number from the list below:", + "median": "Calculate the median number from the list below:", + "expression": "Calculate the numerical expression and provide intermediate results only, for example, for the expression 1 + 3 + 10 - 8, output 4, 14, 6 without displaying the steps.\n\nCalculate the value of the expression below:", + } + inputs = { + "largest number": "You should answer with only one number, no other words. The largest number of the list is: ", + "second largest number": "You should answer with only one number, no other words. The second largest number of the list is: ", + "third largest number": "You should answer with only one number, no other words. The third largest number of the list is: ", + "smallest number": "You should answer with only one number, no other words. The smallest number of the list is: ", + "second smallest number": "You should answer with only one number, no other words. The second smallest number of the list is: ", + "third smallest number": "You should answer with only one number, no other words. The third smallest number of the list is: ", + "median": "You should answer with only one number, no other words. The median number of the list is: ", + "expression": "The value of the numerical expression is: ", + } + for i in range(len(tasks)): + for _ in range(num_samples): + std_out, context = generate_random_list(list_length, min_val, max_val, tasks[i]) + + ret.append({"prompt": prompts[tasks[i]], "context": context, "input": inputs[tasks[i]], "answer": std_out}) + return ret + + +def build_math_find(): + list_length = 60000 # Length of the generated lists + + min_val = 0 # Minimum value for list elements + max_val = 99 # Maximum value for list elements + + ret = generate_math_qa(list_length, min_val, max_val, tasks=["largest number", "second largest number", "third largest number", "smallest number", "second smallest number", "third smallest number", "median"]) + + # Save the data to a JSONL file + fw = jsonlines.open("math_find.jsonl", "w") + fw.write_all(ret) + fw.close() + + +def build_math_calc(): + list_length = 30000 # Length of the generated lists + + min_val = 0 # Minimum value for list elements + max_val = 99 # Maximum value for list elements + + ret = generate_math_qa(list_length, min_val, max_val, tasks=["expression"]) + + # Save the data to a JSONL file + fw = jsonlines.open("math_calc.jsonl", "w") + fw.write_all(ret) + fw.close() + + +def generate_and_store_collections(n, m, min_val, max_val, output_file): + total_elements = n * m + collection = set() + + while len(collection) < total_elements: + collection.add(random.randint(min_val, max_val)) + + collection = list(collection) + random.shuffle(collection) + + collections = [collection[i * m: (i + 1) * m] for i in range(n)] + + with open(output_file, 'w') as file: + json.dump(collections, file) + + +def generate_functions(input_file, min_add, max_add, output_file): + with open(input_file, 'r') as file: + collections = json.load(file) + + function_list = [] + + for i in range(len(collections)): + for t in collections[i]: + function = f"def func_{t}(x):\n" + if i < len(collections) - 1: + next_collection = collections[i + 1] + k = random.choice(next_collection) + addition = random.randint(min_add, max_add) + if addition == 0: + function += f" return func_{k}(x)\n" + elif addition < 0: + function += f" return func_{k}(x) - {-addition}\n" + else: + function += f" return func_{k}(x) + {addition}\n" + else: + addition = random.randint(min_add, max_add) + if addition == 0: + function += f" return x\n" + elif addition < 0: + function += f" return x - {-addition}\n" + else: + function += f" return x + {addition}\n" + function_list.append((f"func_{t}", function)) + + function_list.sort(key=lambda x: int(x[0].split("_")[1])) + + with open(output_file, 'w') as out: + for _, func_text in function_list: + out.write(func_text) + out.write("\n") + + +def generate_code_run_example(collection_file, min_x, max_x, functions_module, functions_file='functions_module.py'): + spec = importlib.util.spec_from_file_location("functions_module", functions_module) + functions = importlib.util.module_from_spec(spec) + spec.loader.exec_module(functions) + # print(functions) + # load all functions in functions_module.py and store them in a string + content = f"\nHere is the content of {functions_file}:\n\n" + with open(functions_module, 'r') as file: + for line in file: + content += line + + with open(collection_file, 'r') as file: + collections = json.load(file) + + + j = random.choice(collections[0]) + x = random.randint(min_x, max_x) + test_sample = { + "context": content, + "answer": getattr(functions, f"func_{j}")(x), + "input": f"Please give me the exact number of the return value of func_{j}({x}). Your response should end with the sentence 'The return value is:'.", + } + + return test_sample + # with jsonlines.open(output_file_samples, mode='w') as writer: + # writer.write_all(test_samples) + # with jsonlines.open(output_file_answers, mode='w') as writer: + # writer.write_all(test_answers) + + + +def build_code_run(): + MAX_NUM_FUNC = 550 + min_val = 1 # minimum value of function indeces + max_val = 2*MAX_NUM_FUNC # maximum value of function indeces + max_add = 17 # maximum value of addition in return expression + min_add = -12 # minimum value of addition in return expression + collections_file = 'collections.json' + functions_file = 'functions_module.py' + #------------------------------------------------------------------------# + # Parameters for generating test samples and answers + num_test = 1 + min_x = -10 + max_x = 10 + n_list = [2, 4, 6, 8, 10] + ret = [] + cnt = -1 + for i in range(len(n_list)): + for _ in range(80): + cnt += 1 + while True: + try: + generate_and_store_collections(n_list[i], int(MAX_NUM_FUNC/n_list[i]), min_val, max_val, collections_file) + + generate_functions(collections_file, min_add, max_add, functions_file) + + example = generate_code_run_example(collections_file, min_x, max_x, functions_file) + example['id'] = cnt + + ret.append(example) + break + except Exception as e: + print(e) + fw = jsonlines.open("code_run.jsonl", 'w') + fw.write_all(ret) + fw.close() + +if __name__ == "__main__": + # os.system("git clone https://github.com/nelson-liu/lost-in-the-middle.git") + # os.system("python3.10 -u lost-in-the-middle/scripts/make_kv_retrieval_data.py --num-keys 3000 --num-examples 500 --output-path kv-retrieval-3000_keys.jsonl.gz") + # os.system("gzip -d kv-retrieval-3000_keys.jsonl.gz") + # build_kv_retrieval() + # build_passkey() + # build_number_string() + # build_math_find() + # build_math_calc() + build_code_run() + diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/functions_module.py b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/functions_module.py new file mode 100644 index 0000000000000000000000000000000000000000..959a06e38f42104f72e68240a82842045bf04343 --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/data/functions_module.py @@ -0,0 +1,1650 @@ +def func_3(x): + return func_490(x) + 9 + +def func_5(x): + return func_147(x) - 5 + +def func_6(x): + return x - 6 + +def func_7(x): + return func_214(x) - 10 + +def func_8(x): + return func_367(x) + 16 + +def func_12(x): + return x - 2 + +def func_13(x): + return func_695(x) - 12 + +def func_14(x): + return x - 9 + +def func_15(x): + return func_28(x) + 12 + +def func_16(x): + return func_400(x) - 11 + +def func_18(x): + return func_516(x) + 4 + +def func_19(x): + return func_361(x) + +def func_21(x): + return func_397(x) - 2 + +def func_22(x): + return func_676(x) - 3 + +def func_23(x): + return func_1099(x) - 9 + +def func_25(x): + return func_287(x) - 4 + +def func_28(x): + return func_772(x) - 1 + +def func_30(x): + return func_242(x) + 9 + +def func_32(x): + return func_132(x) - 3 + +def func_33(x): + return func_674(x) + 12 + +def func_36(x): + return func_288(x) + 5 + +def func_39(x): + return func_990(x) + 9 + +def func_46(x): + return func_761(x) - 9 + +def func_48(x): + return func_965(x) + 12 + +def func_49(x): + return func_320(x) - 12 + +def func_52(x): + return func_441(x) + 9 + +def func_53(x): + return func_911(x) - 9 + +def func_55(x): + return func_825(x) - 2 + +def func_58(x): + return func_387(x) + 17 + +def func_63(x): + return func_650(x) + 5 + +def func_65(x): + return func_1054(x) + +def func_66(x): + return func_659(x) + 4 + +def func_68(x): + return func_928(x) + 12 + +def func_69(x): + return func_923(x) + 8 + +def func_70(x): + return func_25(x) + 6 + +def func_71(x): + return func_39(x) - 7 + +def func_73(x): + return func_880(x) - 6 + +def func_74(x): + return func_25(x) + 6 + +def func_76(x): + return func_740(x) + 6 + +def func_78(x): + return func_137(x) - 3 + +def func_85(x): + return func_911(x) + 4 + +def func_88(x): + return func_963(x) - 7 + +def func_89(x): + return func_116(x) + +def func_90(x): + return func_1049(x) + 3 + +def func_92(x): + return func_706(x) + 12 + +def func_94(x): + return func_979(x) + 10 + +def func_95(x): + return x + 9 + +def func_98(x): + return func_992(x) - 6 + +def func_99(x): + return x - 2 + +def func_101(x): + return func_1080(x) + 10 + +def func_102(x): + return func_565(x) + 15 + +def func_113(x): + return func_309(x) + 17 + +def func_114(x): + return func_625(x) + 7 + +def func_115(x): + return func_1007(x) + 17 + +def func_116(x): + return func_758(x) + 14 + +def func_117(x): + return func_987(x) - 8 + +def func_118(x): + return func_772(x) - 12 + +def func_119(x): + return func_847(x) + 17 + +def func_121(x): + return func_923(x) - 7 + +def func_122(x): + return func_934(x) + 16 + +def func_123(x): + return func_366(x) + 13 + +def func_124(x): + return func_706(x) - 2 + +def func_125(x): + return func_518(x) + 17 + +def func_126(x): + return func_1075(x) - 10 + +def func_130(x): + return func_960(x) - 12 + +def func_131(x): + return func_665(x) + 1 + +def func_132(x): + return func_650(x) + 13 + +def func_134(x): + return func_401(x) + 14 + +def func_137(x): + return func_979(x) - 6 + +def func_140(x): + return func_143(x) - 2 + +def func_141(x): + return func_599(x) - 11 + +def func_143(x): + return x + 3 + +def func_147(x): + return func_954(x) - 6 + +def func_153(x): + return func_371(x) + 3 + +def func_154(x): + return x + 3 + +def func_155(x): + return func_454(x) + +def func_157(x): + return func_126(x) + 13 + +def func_158(x): + return func_319(x) + 10 + +def func_159(x): + return func_510(x) - 12 + +def func_162(x): + return func_707(x) + 8 + +def func_166(x): + return func_802(x) + 1 + +def func_167(x): + return func_1060(x) + 16 + +def func_169(x): + return func_741(x) - 11 + +def func_172(x): + return func_276(x) - 10 + +def func_176(x): + return func_23(x) + 1 + +def func_181(x): + return func_508(x) + 17 + +def func_185(x): + return func_1069(x) - 12 + +def func_188(x): + return func_1016(x) - 6 + +def func_194(x): + return func_661(x) - 1 + +def func_195(x): + return func_892(x) - 9 + +def func_198(x): + return x + 3 + +def func_199(x): + return func_716(x) + 3 + +def func_200(x): + return func_269(x) - 8 + +def func_201(x): + return func_943(x) + 14 + +def func_204(x): + return func_906(x) + 1 + +def func_205(x): + return func_1078(x) - 5 + +def func_207(x): + return func_167(x) - 4 + +def func_208(x): + return func_506(x) - 5 + +def func_209(x): + return func_1019(x) + +def func_214(x): + return x + 9 + +def func_221(x): + return func_903(x) + 3 + +def func_224(x): + return x + 4 + +def func_225(x): + return func_480(x) + 6 + +def func_228(x): + return func_811(x) - 3 + +def func_231(x): + return func_490(x) + 16 + +def func_233(x): + return func_267(x) + 8 + +def func_234(x): + return func_541(x) + 8 + +def func_237(x): + return func_562(x) + +def func_240(x): + return func_225(x) + 4 + +def func_242(x): + return func_432(x) + 8 + +def func_243(x): + return func_627(x) - 5 + +def func_244(x): + return func_23(x) + 9 + +def func_245(x): + return func_567(x) + 16 + +def func_248(x): + return func_115(x) + 5 + +def func_249(x): + return func_158(x) - 4 + +def func_251(x): + return x - 12 + +def func_253(x): + return func_403(x) - 12 + +def func_256(x): + return func_633(x) + 12 + +def func_262(x): + return func_917(x) - 12 + +def func_263(x): + return func_94(x) + 10 + +def func_265(x): + return func_1010(x) + 5 + +def func_267(x): + return func_681(x) + 11 + +def func_268(x): + return func_444(x) - 11 + +def func_269(x): + return func_717(x) + 13 + +def func_272(x): + return func_562(x) - 3 + +def func_274(x): + return func_820(x) + 15 + +def func_275(x): + return func_571(x) - 8 + +def func_276(x): + return x + +def func_277(x): + return func_198(x) - 9 + +def func_278(x): + return func_1095(x) + 16 + +def func_279(x): + return func_525(x) + 3 + +def func_280(x): + return func_1029(x) - 12 + +def func_284(x): + return func_413(x) + 5 + +def func_286(x): + return x - 5 + +def func_287(x): + return func_101(x) - 7 + +def func_288(x): + return func_963(x) + 12 + +def func_289(x): + return func_16(x) + 15 + +def func_291(x): + return func_147(x) + 17 + +def func_292(x): + return func_405(x) + 12 + +def func_294(x): + return func_95(x) + +def func_296(x): + return x + 17 + +def func_297(x): + return func_140(x) + 11 + +def func_299(x): + return func_274(x) + 10 + +def func_301(x): + return func_113(x) + 9 + +def func_302(x): + return func_1086(x) - 9 + +def func_303(x): + return func_521(x) + 17 + +def func_308(x): + return func_727(x) - 11 + +def func_309(x): + return func_302(x) + 5 + +def func_310(x): + return func_48(x) - 12 + +def func_313(x): + return x + 6 + +def func_315(x): + return x - 5 + +def func_316(x): + return func_670(x) + 12 + +def func_317(x): + return func_1005(x) + 15 + +def func_319(x): + return func_98(x) - 4 + +def func_320(x): + return x + 5 + +def func_323(x): + return func_657(x) - 4 + +def func_324(x): + return func_877(x) - 9 + +def func_325(x): + return func_320(x) - 5 + +def func_327(x): + return func_757(x) - 9 + +def func_330(x): + return func_825(x) - 4 + +def func_334(x): + return func_122(x) + +def func_335(x): + return func_445(x) - 7 + +def func_336(x): + return func_153(x) + 16 + +def func_340(x): + return func_758(x) - 10 + +def func_346(x): + return func_85(x) + 1 + +def func_348(x): + return func_567(x) + 8 + +def func_351(x): + return func_22(x) + 5 + +def func_352(x): + return func_527(x) + 16 + +def func_353(x): + return func_860(x) - 7 + +def func_357(x): + return func_878(x) + 1 + +def func_358(x): + return func_960(x) - 11 + +def func_361(x): + return func_48(x) + 5 + +def func_362(x): + return func_134(x) - 2 + +def func_363(x): + return func_1095(x) - 5 + +def func_364(x): + return func_346(x) - 7 + +def func_365(x): + return func_527(x) - 7 + +def func_366(x): + return func_361(x) - 1 + +def func_367(x): + return func_375(x) + 17 + +def func_369(x): + return x - 5 + +def func_370(x): + return func_556(x) + 1 + +def func_371(x): + return func_141(x) - 10 + +def func_374(x): + return x - 2 + +def func_375(x): + return func_828(x) - 6 + +def func_376(x): + return func_251(x) - 5 + +def func_378(x): + return func_231(x) - 8 + +def func_382(x): + return func_1080(x) - 8 + +def func_385(x): + return func_1067(x) + 11 + +def func_386(x): + return func_1003(x) + 14 + +def func_387(x): + return func_98(x) - 9 + +def func_390(x): + return func_1062(x) + 15 + +def func_391(x): + return func_486(x) + 5 + +def func_392(x): + return func_88(x) - 1 + +def func_393(x): + return func_3(x) + 3 + +def func_395(x): + return func_741(x) + +def func_397(x): + return func_730(x) + 17 + +def func_400(x): + return func_253(x) + 1 + +def func_401(x): + return func_376(x) + 10 + +def func_402(x): + return func_556(x) + 9 + +def func_403(x): + return func_506(x) + 13 + +def func_405(x): + return func_572(x) + 13 + +def func_406(x): + return x + 3 + +def func_413(x): + return func_90(x) - 9 + +def func_422(x): + return func_770(x) + 17 + +def func_423(x): + return func_1049(x) - 10 + +def func_428(x): + return func_278(x) + 12 + +def func_429(x): + return func_931(x) - 8 + +def func_432(x): + return func_292(x) - 8 + +def func_434(x): + return x + 2 + +def func_441(x): + return func_297(x) + 11 + +def func_443(x): + return func_696(x) + 12 + +def func_444(x): + return func_124(x) + 16 + +def func_445(x): + return func_618(x) - 5 + +def func_454(x): + return func_113(x) - 4 + +def func_455(x): + return func_325(x) - 2 + +def func_456(x): + return func_1007(x) + 7 + +def func_457(x): + return func_284(x) - 11 + +def func_458(x): + return func_789(x) + 1 + +def func_461(x): + return func_859(x) + 16 + +def func_462(x): + return func_1083(x) - 6 + +def func_463(x): + return func_456(x) + 11 + +def func_466(x): + return func_403(x) - 1 + +def func_469(x): + return func_698(x) + 13 + +def func_470(x): + return func_251(x) + 7 + +def func_473(x): + return func_910(x) + 5 + +def func_474(x): + return func_688(x) + 10 + +def func_480(x): + return func_1007(x) - 7 + +def func_484(x): + return func_673(x) + 3 + +def func_486(x): + return func_12(x) + 2 + +def func_487(x): + return func_70(x) - 11 + +def func_490(x): + return func_455(x) - 2 + +def func_492(x): + return func_53(x) + 7 + +def func_493(x): + return func_288(x) - 8 + +def func_494(x): + return func_757(x) + 4 + +def func_495(x): + return x - 11 + +def func_503(x): + return func_801(x) + 4 + +def func_504(x): + return func_1005(x) - 5 + +def func_505(x): + return func_102(x) - 11 + +def func_506(x): + return func_865(x) + 16 + +def func_508(x): + return func_863(x) + 13 + +def func_510(x): + return func_348(x) - 3 + +def func_514(x): + return func_302(x) - 4 + +def func_516(x): + return func_558(x) + 9 + +def func_518(x): + return func_36(x) + 11 + +def func_521(x): + return func_658(x) + 1 + +def func_523(x): + return func_960(x) - 8 + +def func_525(x): + return func_95(x) + 14 + +def func_526(x): + return func_249(x) - 4 + +def func_527(x): + return x + 8 + +def func_529(x): + return func_627(x) + 17 + +def func_531(x): + return func_323(x) + 14 + +def func_532(x): + return func_1010(x) + 6 + +def func_533(x): + return func_158(x) - 8 + +def func_538(x): + return func_864(x) + 10 + +def func_540(x): + return func_121(x) - 12 + +def func_541(x): + return func_131(x) - 10 + +def func_542(x): + return func_1077(x) + 12 + +def func_543(x): + return func_233(x) + 8 + +def func_545(x): + return func_240(x) + 5 + +def func_547(x): + return func_126(x) + 9 + +def func_548(x): + return x + 6 + +def func_549(x): + return func_395(x) - 8 + +def func_550(x): + return func_650(x) - 5 + +def func_552(x): + return func_324(x) - 5 + +def func_555(x): + return x - 10 + +def func_556(x): + return func_1089(x) + +def func_557(x): + return func_32(x) + 17 + +def func_558(x): + return func_952(x) - 9 + +def func_559(x): + return func_397(x) + 15 + +def func_561(x): + return func_1031(x) + 17 + +def func_562(x): + return func_71(x) - 4 + +def func_564(x): + return func_1095(x) + 4 + +def func_565(x): + return func_432(x) - 7 + +def func_567(x): + return func_778(x) - 5 + +def func_571(x): + return func_552(x) + 2 + +def func_572(x): + return func_251(x) - 8 + +def func_576(x): + return func_251(x) - 1 + +def func_578(x): + return func_860(x) - 12 + +def func_579(x): + return func_141(x) + 16 + +def func_584(x): + return func_249(x) + 16 + +def func_588(x): + return func_1020(x) + 13 + +def func_590(x): + return func_382(x) - 9 + +def func_594(x): + return func_262(x) - 10 + +def func_595(x): + return func_662(x) + 5 + +def func_596(x): + return func_275(x) + 9 + +def func_599(x): + return x + 6 + +def func_600(x): + return func_699(x) + 7 + +def func_610(x): + return x - 1 + +def func_611(x): + return func_169(x) + 3 + +def func_612(x): + return func_979(x) + 6 + +def func_617(x): + return func_875(x) + 7 + +def func_618(x): + return func_313(x) - 2 + +def func_620(x): + return func_796(x) + 9 + +def func_622(x): + return func_1089(x) - 7 + +def func_625(x): + return func_101(x) - 12 + +def func_626(x): + return func_474(x) - 10 + +def func_627(x): + return func_1060(x) - 5 + +def func_633(x): + return func_879(x) - 8 + +def func_634(x): + return func_292(x) + 2 + +def func_637(x): + return func_25(x) + 7 + +def func_638(x): + return func_36(x) - 3 + +def func_639(x): + return func_316(x) + 12 + +def func_641(x): + return func_829(x) - 9 + +def func_644(x): + return func_662(x) - 11 + +def func_645(x): + return func_965(x) + 9 + +def func_647(x): + return func_1007(x) - 10 + +def func_648(x): + return func_548(x) + 1 + +def func_649(x): + return func_692(x) + 13 + +def func_650(x): + return func_1010(x) + +def func_653(x): + return func_1086(x) - 12 + +def func_657(x): + return func_90(x) + 4 + +def func_658(x): + return func_761(x) - 5 + +def func_659(x): + return func_14(x) - 2 + +def func_661(x): + return func_853(x) - 12 + +def func_662(x): + return func_872(x) + 16 + +def func_664(x): + return func_245(x) + 7 + +def func_665(x): + return func_251(x) + 5 + +def func_669(x): + return func_657(x) + 2 + +def func_670(x): + return x + 11 + +def func_672(x): + return x - 4 + +def func_673(x): + return func_869(x) - 4 + +def func_674(x): + return x - 8 + +def func_675(x): + return func_291(x) - 12 + +def func_676(x): + return func_599(x) + 10 + +def func_677(x): + return func_423(x) + 17 + +def func_678(x): + return func_758(x) + 7 + +def func_679(x): + return func_119(x) + 17 + +def func_681(x): + return x - 7 + +def func_683(x): + return func_1029(x) + 3 + +def func_685(x): + return func_248(x) + 11 + +def func_686(x): + return func_1099(x) + 7 + +def func_688(x): + return func_910(x) + 3 + +def func_692(x): + return func_997(x) + 7 + +def func_693(x): + return func_391(x) - 11 + +def func_694(x): + return x + 5 + +def func_695(x): + return func_262(x) + 6 + +def func_696(x): + return func_1013(x) - 5 + +def func_698(x): + return func_890(x) + 5 + +def func_699(x): + return func_965(x) + +def func_701(x): + return func_386(x) + 15 + +def func_702(x): + return func_30(x) + 16 + +def func_703(x): + return func_1007(x) - 6 + +def func_705(x): + return func_964(x) - 1 + +def func_706(x): + return func_308(x) + 14 + +def func_707(x): + return x - 8 + +def func_715(x): + return func_826(x) - 6 + +def func_716(x): + return func_741(x) - 6 + +def func_717(x): + return func_454(x) - 5 + +def func_718(x): + return func_242(x) + +def func_721(x): + return x + 9 + +def func_722(x): + return func_14(x) - 11 + +def func_723(x): + return func_693(x) - 4 + +def func_727(x): + return func_647(x) + 13 + +def func_729(x): + return func_989(x) - 9 + +def func_730(x): + return func_617(x) + 1 + +def func_731(x): + return func_124(x) + 17 + +def func_732(x): + return func_443(x) + 12 + +def func_735(x): + return func_253(x) + 6 + +def func_739(x): + return func_829(x) + +def func_740(x): + return func_369(x) + 12 + +def func_741(x): + return x + 1 + +def func_746(x): + return func_267(x) + 6 + +def func_747(x): + return func_699(x) + 4 + +def func_750(x): + return func_527(x) + 7 + +def func_751(x): + return func_1067(x) + 8 + +def func_754(x): + return func_960(x) + 17 + +def func_756(x): + return x + 14 + +def func_757(x): + return func_58(x) - 5 + +def func_758(x): + return func_1078(x) + 13 + +def func_759(x): + return func_70(x) + 9 + +def func_760(x): + return func_943(x) - 4 + +def func_761(x): + return func_325(x) + 4 + +def func_768(x): + return func_637(x) + +def func_769(x): + return func_692(x) - 9 + +def func_770(x): + return func_679(x) - 12 + +def func_772(x): + return func_1016(x) + +def func_778(x): + return func_224(x) - 11 + +def func_782(x): + return func_118(x) + 4 + +def func_787(x): + return x - 9 + +def func_789(x): + return x + 10 + +def func_792(x): + return x + 4 + +def func_796(x): + return func_770(x) - 7 + +def func_798(x): + return func_1044(x) + 14 + +def func_800(x): + return func_527(x) + 14 + +def func_801(x): + return func_971(x) - 7 + +def func_802(x): + return func_92(x) - 9 + +def func_804(x): + return func_70(x) + 2 + +def func_805(x): + return func_676(x) - 2 + +def func_811(x): + return func_741(x) + 9 + +def func_812(x): + return func_176(x) + 17 + +def func_813(x): + return func_114(x) - 3 + +def func_814(x): + return func_851(x) + 10 + +def func_815(x): + return func_361(x) + 13 + +def func_819(x): + return func_730(x) + 9 + +def func_820(x): + return func_248(x) - 11 + +def func_821(x): + return func_233(x) - 10 + +def func_823(x): + return func_819(x) - 3 + +def func_824(x): + return func_622(x) + 5 + +def func_825(x): + return func_176(x) + 15 + +def func_826(x): + return func_1047(x) - 5 + +def func_827(x): + return func_625(x) + 3 + +def func_828(x): + return func_126(x) - 10 + +def func_829(x): + return func_815(x) + 12 + +def func_830(x): + return func_863(x) + 3 + +def func_832(x): + return func_401(x) - 11 + +def func_836(x): + return func_492(x) + 12 + +def func_838(x): + return func_153(x) + 14 + +def func_840(x): + return x - 3 + +def func_842(x): + return func_253(x) - 3 + +def func_843(x): + return func_987(x) + 1 + +def func_845(x): + return func_463(x) - 7 + +def func_846(x): + return func_678(x) + 3 + +def func_847(x): + return func_199(x) - 6 + +def func_851(x): + return func_505(x) - 4 + +def func_853(x): + return func_990(x) + 8 + +def func_856(x): + return func_397(x) + 16 + +def func_857(x): + return func_579(x) - 3 + +def func_859(x): + return func_406(x) + 1 + +def func_860(x): + return func_378(x) + 14 + +def func_861(x): + return func_958(x) + +def func_863(x): + return func_361(x) - 4 + +def func_864(x): + return func_730(x) + 2 + +def func_865(x): + return x - 6 + +def func_866(x): + return x + 4 + +def func_869(x): + return func_369(x) + 1 + +def func_871(x): + return func_265(x) + 3 + +def func_872(x): + return func_902(x) + 17 + +def func_873(x): + return func_1076(x) + 14 + +def func_875(x): + return func_309(x) + 1 + +def func_877(x): + return func_750(x) + 9 + +def func_878(x): + return func_1021(x) - 11 + +def func_879(x): + return func_423(x) + 16 + +def func_880(x): + return func_1042(x) + 7 + +def func_882(x): + return func_527(x) - 1 + +def func_886(x): + return func_1091(x) + +def func_887(x): + return func_208(x) + 12 + +def func_889(x): + return func_36(x) - 11 + +def func_890(x): + return func_1091(x) - 8 + +def func_891(x): + return func_492(x) + 14 + +def func_892(x): + return func_233(x) + 16 + +def func_896(x): + return func_827(x) + 7 + +def func_901(x): + return func_284(x) + 11 + +def func_902(x): + return func_406(x) + 5 + +def func_903(x): + return func_23(x) + 2 + +def func_906(x): + return func_301(x) - 1 + +def func_907(x): + return func_578(x) + 2 + +def func_910(x): + return func_195(x) - 9 + +def func_911(x): + return func_983(x) + 7 + +def func_912(x): + return x + 15 + +def func_913(x): + return x - 6 + +def func_915(x): + return func_1080(x) - 2 + +def func_917(x): + return func_693(x) - 7 + +def func_920(x): + return func_516(x) + 16 + +def func_923(x): + return func_336(x) - 1 + +def func_924(x): + return func_443(x) - 12 + +def func_927(x): + return func_7(x) + 15 + +def func_928(x): + return func_335(x) + 2 + +def func_931(x): + return func_245(x) + +def func_934(x): + return func_1042(x) - 1 + +def func_936(x): + return func_137(x) + 6 + +def func_937(x): + return func_915(x) + 4 + +def func_939(x): + return func_353(x) + 14 + +def func_940(x): + return func_757(x) - 7 + +def func_943(x): + return func_208(x) + 14 + +def func_945(x): + return func_330(x) + 5 + +def func_948(x): + return func_686(x) - 11 + +def func_949(x): + return func_757(x) + 13 + +def func_950(x): + return x + 5 + +def func_952(x): + return func_493(x) + 13 + +def func_953(x): + return x + 17 + +def func_954(x): + return x - 7 + +def func_955(x): + return func_772(x) + 2 + +def func_957(x): + return func_948(x) + +def func_958(x): + return func_578(x) - 10 + +def func_960(x): + return func_677(x) - 6 + +def func_962(x): + return func_564(x) + 11 + +def func_963(x): + return func_1007(x) - 5 + +def func_964(x): + return func_286(x) + 9 + +def func_965(x): + return func_375(x) + 7 + +def func_971(x): + return func_953(x) - 10 + +def func_972(x): + return func_564(x) - 12 + +def func_973(x): + return x + 11 + +def func_974(x): + return func_637(x) + 3 + +def func_976(x): + return func_696(x) - 6 + +def func_978(x): + return func_461(x) - 4 + +def func_979(x): + return func_672(x) - 9 + +def func_983(x): + return func_648(x) + 4 + +def func_985(x): + return func_564(x) - 10 + +def func_986(x): + return func_936(x) - 5 + +def func_987(x): + return func_873(x) + 3 + +def func_988(x): + return x + 7 + +def func_989(x): + return func_335(x) + 8 + +def func_990(x): + return func_674(x) - 9 + +def func_991(x): + return func_1067(x) + 1 + +def func_992(x): + return func_351(x) + +def func_993(x): + return func_1043(x) + 7 + +def func_996(x): + return func_896(x) + 13 + +def func_997(x): + return func_688(x) - 6 + +def func_1000(x): + return func_986(x) + 5 + +def func_1003(x): + return func_296(x) - 6 + +def func_1004(x): + return func_463(x) - 1 + +def func_1005(x): + return func_92(x) + 1 + +def func_1007(x): + return func_572(x) - 1 + +def func_1008(x): + return func_367(x) + 17 + +def func_1010(x): + return func_224(x) - 12 + +def func_1013(x): + return func_262(x) + 15 + +def func_1016(x): + return func_276(x) + 1 + +def func_1019(x): + return x - 10 + +def func_1020(x): + return func_782(x) + 8 + +def func_1021(x): + return x + 12 + +def func_1027(x): + return func_405(x) + 2 + +def func_1029(x): + return func_221(x) + 3 + +def func_1030(x): + return func_237(x) - 8 + +def func_1031(x): + return func_12(x) - 2 + +def func_1032(x): + return func_813(x) + 16 + +def func_1035(x): + return func_294(x) + 5 + +def func_1037(x): + return func_954(x) + 17 + +def func_1042(x): + return func_23(x) + 11 + +def func_1043(x): + return func_845(x) + 6 + +def func_1044(x): + return x - 7 + +def func_1045(x): + return x + 11 + +def func_1047(x): + return func_288(x) + 1 + +def func_1049(x): + return func_88(x) - 6 + +def func_1051(x): + return func_63(x) - 4 + +def func_1053(x): + return func_832(x) - 5 + +def func_1054(x): + return func_761(x) - 3 + +def func_1059(x): + return func_397(x) + 12 + +def func_1060(x): + return func_600(x) + 17 + +def func_1061(x): + return func_826(x) + 6 + +def func_1062(x): + return func_549(x) + 4 + +def func_1067(x): + return func_963(x) + 2 + +def func_1069(x): + return func_541(x) + 7 + +def func_1075(x): + return x + 7 + +def func_1076(x): + return func_845(x) + 11 + +def func_1077(x): + return func_661(x) - 10 + +def func_1078(x): + return func_634(x) - 7 + +def func_1079(x): + return func_928(x) - 11 + +def func_1080(x): + return func_658(x) + 6 + +def func_1082(x): + return x + 6 + +def func_1083(x): + return func_237(x) + 4 + +def func_1086(x): + return func_1082(x) - 3 + +def func_1089(x): + return func_625(x) + 14 + +def func_1090(x): + return func_760(x) - 10 + +def func_1091(x): + return func_393(x) + 13 + +def func_1093(x): + return func_244(x) - 5 + +def func_1094(x): + return func_813(x) - 9 + +def func_1095(x): + return func_387(x) - 8 + +def func_1096(x): + return func_185(x) - 8 + +def func_1098(x): + return func_873(x) + 1 + +def func_1099(x): + return func_456(x) - 8 + +def func_1100(x): + return func_692(x) + diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/requirements.txt b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..a28f5f9485c8bc31498d9e69c69d54d285debcf1 --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/requirements.txt @@ -0,0 +1,9 @@ +openai +tiktoken +rouge +torch +transformers +accelerate +evaluate +xopen +python-dotenv \ No newline at end of file diff --git a/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/scripts/download_dataset.sh b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/scripts/download_dataset.sh new file mode 100644 index 0000000000000000000000000000000000000000..6e626d79056708b8bf1f9f2425098649649d5813 --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/InfiniteBench/scripts/download_dataset.sh @@ -0,0 +1,6 @@ +#!/bin/bash +save_dir=data +mkdir ${save_dir} +for file in code_debug code_run kv_retrieval longbook_choice_eng longbook_qa_chn longbook_qa_eng longbook_sum_eng longdialogue_qa_eng math_calc math_find number_string passkey; do + wget -c https://huggingface.co/datasets/xinrongzhang2022/InfiniteBench/resolve/main/${file}.jsonl?download=true -O ./${save_dir}/${file}.jsonl +done \ No newline at end of file diff --git a/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/MMLU_Pro_rewritten.py b/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/MMLU_Pro_rewritten.py new file mode 100644 index 0000000000000000000000000000000000000000..acc3e2a811f0ed61fa31fb4f99efe31c61b1e455 --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/MMLU_Pro_rewritten.py @@ -0,0 +1,341 @@ +# MMLU_Pro_rewritten.py +# Description: Script to perform MMLU-Pro benchmarking +# +#################################################################################################################### +# Imports +import os +import threading +import time +import toml +from tqdm import tqdm +from concurrent.futures import ThreadPoolExecutor +import logging +from openai import OpenAI +from datasets import load_dataset +import json +import re +# +################################################################################################################## +# +# Functions: + + +# Set up logging +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') +logger = logging.getLogger(__name__) + + +def load_mmlu_pro_config(**kwargs): + # Get the directory of the current script + script_dir = os.path.dirname(os.path.abspath(__file__)) + + # Construct the full path to config.toml + config_path = os.path.join(script_dir, 'config.toml') + + # Load the config + config = toml.load(config_path) + + # Update config with provided kwargs + for key, value in kwargs.items(): + if key in config["server"]: + config["server"][key] = value + elif key in config["test"]: + config["test"][key] = value + elif key in config["log"]: + config["log"][key] = value + + return config + +# client_initializer.py +def initialize_client(config): + try: + return OpenAI( + base_url=config["server"]["url"], + api_key=config["server"]["api_key"], + timeout=config["server"]["timeout"] + ) + except Exception as e: + logger.error(f"Failed to initialize OpenAI client: {e}") + raise + +# dataset_loader.py +def load_mmlu_pro(): + try: + dataset = load_dataset("TIGER-Lab/MMLU-Pro") + test_df, val_df = dataset["test"], dataset["validation"] + return preprocess(test_df), preprocess(val_df) + except Exception as e: + logger.error(f"Error loading MMLU-Pro dataset: {e}") + raise + +def preprocess(data): + res = {} + for item in data: + options = [opt for opt in item["options"] if opt != "N/A"] + item["options"] = options + category = item["category"] + if category not in res: + res[category] = [] + res[category].append(item) + return res + +# prompt_creator.py +def create_prompt(cot_examples, question, options, config): + style = config["inference"]["style"] + system_prompt = config["inference"]["system_prompt"] + + def format_example(q, opts, cot=""): + if not cot: + cot = "Let's think step by step." + cot = cot[3:] if cot.startswith("A: ") else cot + example = f"Question: {q}\nOptions: " + example += "\n".join(f"{chr(65 + i)}. {opt}" for i, opt in enumerate(opts)) + return example.strip(), cot.strip() + + if style == "multi_chat": + messages = [{"role": "system", "content": system_prompt}] + for ex in cot_examples: + ex_text, cot = format_example(ex["question"], ex["options"], ex["cot_content"]) + messages.extend([ + {"role": "user", "content": ex_text}, + {"role": "assistant", "content": f"Answer: {cot}"} + ]) + q_text, _ = format_example(question, options) + messages.append({"role": "user", "content": q_text}) + return messages + elif style == "single_chat": + prompt = f"{system_prompt}\n\n" + for ex in cot_examples: + ex_text, cot = format_example(ex["question"], ex["options"], ex["cot_content"]) + prompt += f"{ex_text}\nAnswer: {cot}\n\n" + q_text, _ = format_example(question, options) + prompt += f"{q_text}\nAnswer: Let's think step by step." + return [{"role": "user", "content": prompt}] + else: # no_chat + prompt = f"{system_prompt}\n\n" + for ex in cot_examples: + ex_text, cot = format_example(ex["question"], ex["options"], ex["cot_content"]) + prompt += f"{ex_text}\nAnswer: {cot}\n\n" + q_text, _ = format_example(question, options) + prompt += f"{q_text}\nAnswer: Let's think step by step." + return prompt + +# answer_extractor.py +def extract_answer(text): + patterns = [ + r"answer is \(?([A-J])\)?", + r".*[aA]nswer:\s*\(?([A-J])\)?", + r"\b([A-J])\b(?!.*\b[A-J]\b)" + ] + + for pattern in patterns: + match = re.search(pattern, text, re.IGNORECASE | re.DOTALL) + if match: + return match.group(1).upper() + + logger.warning(f"Failed to extract answer from: {text}") + return None + +# question_evaluator.py +def run_single_question(question, cot_examples, client, config): + max_retries = 3 + for attempt in range(max_retries): + try: + prompt = create_prompt(cot_examples, question['question'], question['options'], config) + + if config["inference"]["style"] == "no_chat": + response = client.completions.create( + model=config["server"]["model"], + prompt=prompt, + temperature=config["inference"]["temperature"], + max_tokens=config["inference"]["max_tokens"], + top_p=config["inference"]["top_p"], + frequency_penalty=0, + presence_penalty=0, + stop=["Question:"], + timeout=config["server"]["timeout"], + ) + response_text = response.choices[0].text.strip() + else: + response = client.chat.completions.create( + model=config["server"]["model"], + messages=prompt, + temperature=config["inference"]["temperature"], + max_tokens=config["inference"]["max_tokens"], + top_p=config["inference"]["top_p"], + frequency_penalty=0, + presence_penalty=0, + stop=["Question:"], + timeout=config["server"]["timeout"], + ) + response_text = response.choices[0].message.content.strip() + + pred = extract_answer(response_text) + usage = response.usage + + return prompt, response_text, pred, usage + + except Exception as e: + logger.warning(f"Attempt {attempt + 1} failed: {e}") + if attempt == max_retries - 1: + logger.error(f"All attempts failed for question: {question['question_id']}") + return None, None, None, None + time.sleep(3) # Wait before retrying + +# result_processor.py +def save_results(results, output_path, lock): + max_retries = 3 + for attempt in range(max_retries): + try: + with lock: + with open(output_path, 'w') as f: + json.dump(results, f, indent=2) + return + except Exception as e: + logger.warning(f"Attempt {attempt + 1} to save results failed: {e}") + if attempt == max_retries - 1: + logger.error(f"Failed to save results to {output_path}") + time.sleep(1) # Wait before retrying + +def save_summary(category_record, output_path, lock): + max_retries = 3 + for attempt in range(max_retries): + try: + with lock: + with open(output_path, 'w') as f: + json.dump(category_record, f, indent=2) + return + except Exception as e: + logger.warning(f"Attempt {attempt + 1} to save summary failed: {e}") + if attempt == max_retries - 1: + logger.error(f"Failed to save summary to {output_path}") + time.sleep(1) # Wait before retrying + +def update_results(results, category_record, question, pred, answer): + category = question['category'] + + if category not in category_record: + category_record[category] = {"correct": 0, "total": 0} + + category_record[category]["total"] += 1 + if pred == answer: + category_record[category]["correct"] += 1 + + result = { + "question_id": question['question_id'], + "category": category, + "question": question['question'], + "options": question['options'], + "pred": pred, + "answer": answer, + "correct": pred == answer + } + results.append(result) + + return results, category_record + +def process_and_save_results(question, pred, client, config, results, category_record, output_dir, lock): + results, category_record = update_results(results, category_record, question, pred, question['answer']) + + output_res_path = os.path.join(output_dir, f"{question['category']}_result.json") + output_summary_path = os.path.join(output_dir, f"{question['category']}_summary.json") + + save_results(results, output_res_path, lock) + save_summary(category_record, output_summary_path, lock) + + return results, category_record + +def generate_final_report(category_record, output_dir): + total_correct = sum(cat["correct"] for cat in category_record.values()) + total_questions = sum(cat["total"] for cat in category_record.values()) + overall_accuracy = total_correct / total_questions if total_questions > 0 else 0 + + report = f"MMLU-Pro Benchmark Final Report\n" + report += f"================================\n\n" + report += f"Overall Accuracy: {overall_accuracy:.2%} ({total_correct}/{total_questions})\n\n" + report += f"Category Breakdown:\n" + for category, stats in category_record.items(): + accuracy = stats["correct"] / stats["total"] if stats["total"] > 0 else 0 + report += f" {category}: {accuracy:.2%} ({stats['correct']}/{stats['total']})\n" + + report_path = os.path.join(output_dir, "final_report.txt") + with open(report_path, 'w') as f: + f.write(report) + + logger.info(f"Final report saved to {report_path}") + +def mmlu_pro_main(): + # Load configuration + config = load_mmlu_pro_config() + + # Initialize OpenAI client + client = initialize_client(config) + + # Load and preprocess the MMLU-Pro dataset + test_data, dev_data = load_mmlu_pro() + if test_data is None or dev_data is None: + logger.error("Failed to load dataset. Exiting.") + return + + # Prepare output directory + output_dir = os.path.join("eval_results", config["server"]["model"].replace("/", "-")) + os.makedirs(output_dir, exist_ok=True) + + # Initialize results storage + results = [] + category_record = {} + lock = threading.Lock() + + # Set a failure threshold to cancel the benchmark if too many questions fail + max_failed_questions = 6 + failed_questions = 0 + + # Process each subject + for subject, questions in test_data.items(): + logger.info(f"Processing subject: {subject}") + cot_examples = dev_data[subject] + + # Use ThreadPoolExecutor for parallel processing + with ThreadPoolExecutor(max_workers=config["test"]["parallel"]) as executor: + futures = [] + for question in questions: + future = executor.submit(run_single_question, question, cot_examples, client, config) + futures.append((future, question)) + + # Process results as they complete + for future, question in tqdm(futures, total=len(futures)): + prompt, response, pred, usage = future.result() + + # Check if the question failed and increment the failure count + if pred is None: + failed_questions += 1 + logger.warning(f"Failed question count: {failed_questions}/{max_failed_questions}") + + # Stop the entire process if too many questions fail + if failed_questions >= max_failed_questions: + logger.error(f"Too many failed questions. Stopping the benchmark for {subject}.") + return + + # Process and save results if the question was answered + if pred is not None: + results, category_record = process_and_save_results( + question, pred, client, config, results, category_record, output_dir, lock + ) + + # Save final results for the subject + save_results(results, os.path.join(output_dir, f"{subject}_final_result.json"), lock) + save_summary(category_record, os.path.join(output_dir, f"{subject}_final_summary.json"), lock) + + # Generate and save final report + generate_final_report(category_record, output_dir) + + logger.info(f"Evaluation complete. Results saved in {output_dir}") + +def run_mmlu_pro_benchmark(): + start_time = time.time() + mmlu_pro_main() + end_time = time.time() + logger.info(f"Total execution time: {end_time - start_time:.2f} seconds") +# +# End of file +#################################################################################################### diff --git a/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/__init__.py b/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/config.toml b/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/config.toml new file mode 100644 index 0000000000000000000000000000000000000000..632fdf1a0d41e61c017a4e17d2356c8d9a40389a --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/config.toml @@ -0,0 +1,30 @@ +# Comment to be included in the beginning of the final report. +comment = "" + +[server] +url = "http://localhost:11434/v1" +api_key = "api key" +model = "llama3" +timeout = 600.0 + +[inference] +# Ssettings below are from evaluate_from_local.py for VLLM on TIGER-AI-Lab/MMLU-Pro +temperature = 0.0 +top_p = 1.0 # not specified but default for VLLM +max_tokens = 2048 +# The variable {subject} will be replaced with appropriate value in runtime. +system_prompt = "The following are multiple choice questions (with answers) about {subject}. Think step by step and then finish your answer with \"the answer is (X)\" where X is the correct letter choice." +# "multi_chat" inserts COT examples into multi-turn messages. Use for instruct/chat models. +# "no_chat" uses v1/completion api. Use for non-instruct/chat model. +# "single_chat" (from the script for GPT-4O) inserts all the COT examples and question into a single message. Not recommended, use only for legacy compatibility. +style = "multi_chat" + +[test] +categories = ['biology', 'business', 'chemistry', 'computer science', 'economics', 'engineering', 'health', 'history', 'law', 'math', 'philosophy', 'physics', 'psychology', 'other'] +parallel = 1 + +[log] +# Verbosity between 0-2 +verbosity = 0 +# If true, logs exact prompt sent to the model in the test result files. +log_prompt = true diff --git a/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/mmlu_pro_test.py b/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/mmlu_pro_test.py new file mode 100644 index 0000000000000000000000000000000000000000..bc9b1dff9efd36954ce7a7f2d6ffc5eb90d61580 --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/mmlu_pro_test.py @@ -0,0 +1,232 @@ +# Test the load_config function +def test_load_config(): + import sys + original_argv = sys.argv + #sys.argv = ["run_openai.py", "-c", "test_config.toml", "-u", "http://test.com", "-m", "test-model"] + + config = load_config() + + assert config["server"]["url"] == "http://test.com" + assert config["server"]["model"] == "test-model" + + sys.argv = original_argv + print("load_config test passed") + +def test_load_mmlu_pro(): + test_df, val_df = load_mmlu_pro() + assert test_df is not None + assert val_df is not None + assert isinstance(test_df, dict) + assert isinstance(val_df, dict) + print("load_mmlu_pro test passed") + + +def test_initialize_client(): + test_config = { + "server": { + "url": "http://test.com", + "api_key": "test_key", + "timeout": 30 + } + } + + client = initialize_client(test_config) + + assert client.base_url == "http://test.com" + assert client.api_key == "test_key" + assert client.timeout == 30 + + print("initialize_client test passed") + + +test_initialize_client() + +def test_preprocess(): + sample_data = [ + {"category": "math", "options": ["A", "B", "N/A", "C"]}, + {"category": "science", "options": ["X", "Y", "Z"]} + ] + processed = preprocess(sample_data) + assert "math" in processed + assert "science" in processed + assert len(processed["math"][0]["options"]) == 3 + assert "N/A" not in processed["math"][0]["options"] + assert len(processed["science"][0]["options"]) == 3 + print("preprocess test passed") + +test_load_mmlu_pro() +test_preprocess() + + +test_load_config() + + +def test_create_prompt(): + config = { + "inference": { + "style": "multi_chat", + "system_prompt": "You are a helpful assistant." + } + } + cot_examples = [{ + "question": "What is 2+2?", + "options": ["3", "4", "5"], + "cot_content": "Let's add 2 and 2. 2+2 = 4." + }] + question = "What is 3+3?" + options = ["5", "6", "7"] + + # Test multi_chat + result = create_prompt(cot_examples, question, options, config) + assert isinstance(result, list) + assert len(result) == 4 + assert result[0]["role"] == "system" + assert result[-1]["role"] == "user" + + # Test single_chat + config["inference"]["style"] = "single_chat" + result = create_prompt(cot_examples, question, options, config) + assert isinstance(result, list) + assert len(result) == 1 + assert result[0]["role"] == "user" + + # Test no_chat + config["inference"]["style"] = "no_chat" + result = create_prompt(cot_examples, question, options, config) + assert isinstance(result, str) + assert "What is 2+2?" in result + assert "What is 3+3?" in result + + print("create_prompt test passed") + +test_create_prompt() + + +def test_extract_answer(): + test_cases = [ + ("The answer is (B)", "B"), + ("After careful consideration, I believe the answer is C.", "C"), + ( + "Let's analyze each option:\nA. Incorrect\nB. Incorrect\nC. Correct\nD. Incorrect\nTherefore, the answer is C.", + "C"), + ("A. GHTIS\nB. MCU\nC. UBT\nD. ALIN\n\nThe correct answer is B. MCU.", "B"), + ("There is no clear answer in this text.", None), + ("The options are A, B, C, and D. I think B is the best answer.", "B") + ] + + for text, expected in test_cases: + result = extract_answer(text) + assert result == expected, f"Failed on input '{text}'. Expected {expected}, got {result}" + + print("extract_answer test passed") + + +test_extract_answer() + +from unittest.mock import Mock + +def test_run_single_question(): + # Mock OpenAI client + mock_client = Mock() + mock_response = Mock() + mock_response.choices = [Mock(text="The answer is B", message=Mock(content="The answer is B"))] + mock_response.usage = Mock(prompt_tokens=10, completion_tokens=20, total_tokens=30) + mock_client.completions.create.return_value = mock_response + mock_client.chat.completions.create.return_value = mock_response + + # Mock configuration + config = { + "inference": { + "style": "no_chat", + "system_prompt": "You are a helpful assistant.", + "temperature": 0.7, + "max_tokens": 100, + "top_p": 1.0 + }, + "server": { + "model": "test-model", + "timeout": 30 + } + } + + # Mock question and examples + question = { + "question": "What is 2+2?", + "options": ["3", "4", "5"] + } + cot_examples = [] + + # Test no_chat style + prompt, response, pred, usage = run_single_question(question, cot_examples, mock_client, config) + assert prompt is not None + assert response == "The answer is B" + assert pred == "B" + assert usage.prompt_tokens == 10 + assert usage.completion_tokens == 20 + assert usage.total_tokens == 30 + + # Test chat style + config["inference"]["style"] = "multi_chat" + prompt, response, pred, usage = run_single_question(question, cot_examples, mock_client, config) + assert prompt is not None + assert response == "The answer is B" + assert pred == "B" + assert usage.prompt_tokens == 10 + assert usage.completion_tokens == 20 + assert usage.total_tokens == 30 + + print("run_single_question test passed") + +test_run_single_question() + + +def test_save_and_update_functions(): + # Create a temporary directory for test files + with tempfile.TemporaryDirectory() as tmpdir: + lock = threading.Lock() + results = [] + category_record = {} + + # Test question + question = { + 'question_id': '1', + 'category': 'math', + 'question': 'What is 2+2?', + 'options': ['3', '4', '5'], + 'answer': 'B' + } + + # Test update_results + results, category_record = update_results(results, category_record, question, 'B', 'B') + assert len(results) == 1 + assert category_record['math']['correct'] == 1 + assert category_record['math']['total'] == 1 + + # Test save_results and save_summary + results_path = os.path.join(tmpdir, 'results.json') + summary_path = os.path.join(tmpdir, 'summary.json') + + save_results(results, results_path, lock) + save_summary(category_record, summary_path, lock) + + assert os.path.exists(results_path) + assert os.path.exists(summary_path) + + # Test process_and_save_results + config = {'server': {'model': 'test-model'}} + client = None # We don't need a real client for this test + + results, category_record = process_and_save_results(question, 'B', client, config, results, category_record, + tmpdir, lock) + + assert len(results) == 2 + assert category_record['math']['correct'] == 2 + assert category_record['math']['total'] == 2 + + assert os.path.exists(os.path.join(tmpdir, 'math_result.json')) + assert os.path.exists(os.path.join(tmpdir, 'math_summary.json')) + + print("save_and_update_functions tests passed") + + +test_save_and_update_functions() \ No newline at end of file diff --git a/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/run_openai.py b/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/run_openai.py new file mode 100644 index 0000000000000000000000000000000000000000..4348ca6aa9fccca8d1fdbe0213880f6fbcd812b7 --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/MMLU_Pro/run_openai.py @@ -0,0 +1,546 @@ +# Script taken from: https://github.com/chigkim/Ollama-MMLU-Pro +# No changes made +import os +import re +import json +import time +import random +from tqdm import tqdm +from openai import OpenAI +from datasets import load_dataset +from concurrent.futures import ThreadPoolExecutor, as_completed +import threading +from datetime import datetime, timedelta +import codecs +import toml +import argparse +import queue +import numpy as np +import copy + +parser = argparse.ArgumentParser( + prog="python3 run_openai.py", + description="Run MMLU Pro Benchmark for a local LLM via OpenAI Compatible API.", + epilog="Specify options above to override one or more settings from config.", +) +parser.add_argument( + "-c", + "--config", + help="Configuration file. Default=config.toml", + default="config.toml", +) +parser.add_argument( + "-u", + "--url", + help="server url", +) +parser.add_argument("-a", "--api", help="api key") +parser.add_argument("-m", "--model", help="Model name") +parser.add_argument( + "--timeout", + type=float, + help="Request timeout in seconds", +) +parser.add_argument("--category", type=str) +parser.add_argument("-p", "--parallel", type=int, help="Number of parallel requests") +parser.add_argument("-v", "--verbosity", type=int, help="Verbosity level 0-2") +parser.add_argument( + "--log_prompt", + help="Writes exact prompt and response into log.txt", + action="store_true", +) +parser.add_argument( + "--comment", type=str, help="Comment to be included in the final report." +) +args = parser.parse_args() +config = toml.load(open(args.config)) +if args.url: + config["server"]["url"] = args.url +if args.api: + config["server"]["api_key"] = args.api +if args.model: + config["server"]["model"] = args.model +if args.timeout: + config["server"]["timeout"] = args.timeout +if args.category: + config["test"]["categories"] = [args.category] +if args.parallel: + config["test"]["parallel"] = args.parallel +if args.verbosity: + config["log"]["verbosity"] = args.verbosity +if args.log_prompt: + config["log"]["log_prompt"] = args.log_prompt +if args.comment: + config["comment"] = args.comment + + +client = OpenAI( + base_url=config["server"]["url"], + api_key=config["server"]["api_key"], + timeout=config["server"]["timeout"], +) + + +def log(message): + print(message) + with codecs.open(log_path, "a", "utf-8") as file: + file.write(message + "\n") + + +def get_chat_completion(messages): + try: + response = client.chat.completions.create( + model=config["server"]["model"], + messages=messages, + temperature=config["inference"]["temperature"], + max_tokens=config["inference"]["max_tokens"], + top_p=config["inference"]["top_p"], + frequency_penalty=0, + presence_penalty=0, + stop=["Question:"], + timeout=config["server"]["timeout"], + ) + try: + usage_q.put( + (response.usage.prompt_tokens, response.usage.completion_tokens) + ) + except: + pass + return response.choices[0].message.content.strip() + except Exception as e: + print("Resubmitting, Error: ", e) + time.sleep(3) + return get_chat_completion(messages) + + +def get_completion(prompt): + try: + response = client.completions.create( + model=config["server"]["model"], + prompt=prompt, + temperature=config["inference"]["temperature"], + max_tokens=config["inference"]["max_tokens"], + top_p=config["inference"]["top_p"], + frequency_penalty=0, + presence_penalty=0, + stop=["Question:"], + timeout=config["server"]["timeout"], + ) + try: + usage_q.put( + (response.usage.prompt_tokens, response.usage.completion_tokens) + ) + except: + pass + if response.choices: + return response.choices[0].text.strip() + elif response.content: + return response.content.strip() + print("Can't get response.") + return None + except Exception as e: + print("Resubmitting, Error: ", e) + time.sleep(3) + return get_completion(prompt) + + +def load_mmlu_pro(): + dataset = load_dataset("TIGER-Lab/MMLU-Pro") + test_df, val_df = dataset["test"], dataset["validation"] + test_df = preprocess(test_df) + val_df = preprocess(val_df) + return test_df, val_df + + +def preprocess(test_df): + res_df = [] + for each in test_df: + options = [] + for opt in each["options"]: + if opt == "N/A": + continue + options.append(opt) + each["options"] = options + res_df.append(each) + res = {} + for each in res_df: + if each["category"] not in res: + res[each["category"]] = [] + res[each["category"]].append(each) + return res + + +def format_example(question, options, cot_content=""): + if cot_content == "": + cot_content = "Let's think step by step." + if cot_content.startswith("A: "): + cot_content = cot_content[3:] + example = "Question: {}\nOptions: ".format(question) + choice_map = "ABCDEFGHIJ" + for i, opt in enumerate(options): + example += "{}. {}\n".format(choice_map[i], opt) + return example.strip(), cot_content.strip() + + +def multi_chat_prompt(cot_examples, question, options): + messages = [ + { + "role": "system", + "content": config["inference"]["system_prompt"], + }, + ] + for each in cot_examples: + example, cot_content = format_example( + each["question"], each["options"], each["cot_content"] + ) + messages.append({"role": "user", "content": example}) + messages.append({"role": "assistant", "content": "Answer: " + cot_content}) + example, cot_content = format_example(question, options) + messages.append({"role": "user", "content": example}) + return messages + + +def single_chat_prompt(cot_examples, question, options): + messages = [ + { + "role": "system", + "content": config["inference"]["system_prompt"], + }, + ] + prompt = no_chat_prompt(cot_examples, question, options, no_system=True) + messages.append({"role": "user", "content": prompt}) + return messages + + +def no_chat_prompt(cot_examples, question, options, no_system=False): + prompt = config["inference"]["system_prompt"] + "\n\n" + if no_system: + prompt = "" + for each in cot_examples: + example, cot_content = format_example( + each["question"], each["options"], each["cot_content"] + ) + prompt += example + "\n" + prompt += "Answer: " + cot_content + "\n\n" + example, cot_content = format_example(question, options) + prompt += example + "\n" + prompt += "Answer: " + cot_content + return prompt + + +def extract_answer(text): + pattern = r"answer is \(?([ABCDEFGHIJ])\)?" + match = re.search(pattern, text) + if match: + return match.group(1) + else: + return extract_again(text) + + +def extract_again(text): + pattern = r".*[aA]nswer:\s*\(?([A-J])\)?" + match = re.search(pattern, text) + if match: + return match.group(1) + else: + return extract_final(text) + + +def extract_final(text): + pattern = r"\b[A-J]\b(?!.*\b[A-J]\b)" + match = re.search(pattern, text, re.DOTALL) + if match: + return match[0] + else: + if config["log"]["verbosity"] >= 1: + print("Extraction failed:\n", text) + return None + + +def run_single_question(single_question, cot_examples_dict, exist_result): + exist = True + q_id = single_question["question_id"] + for each in exist_result: + if ( + q_id == each["question_id"] + and single_question["question"] == each["question"] + ): + if config["log"]["verbosity"] >= 1: + print("already exists, skipping.") + return None, None, None, exist + exist = False + category = single_question["category"] + cot_examples = cot_examples_dict[category] + question = single_question["question"] + options = single_question["options"] + try: + if config["inference"]["style"] == "single_chat": + prompt = single_chat_prompt(cot_examples, question, options) + response = get_chat_completion(prompt) + elif config["inference"]["style"] == "multi_chat": + prompt = multi_chat_prompt(cot_examples, question, options) + response = get_chat_completion(prompt) + elif config["inference"]["style"] == "no_chat": + prompt = no_chat_prompt(cot_examples, question, options) + response = get_completion(prompt) + except Exception as e: + print("error", e) + return None, None, None, exist + pred = extract_answer(response) + return prompt, response, pred, exist + + +def update_result(output_res_path, lock): + category_record = {} + res = [] + success = False + while not success: + try: + if os.path.exists(output_res_path): + with lock: + with open(output_res_path, "r") as fi: + res = json.load(fi) + for each in res: + category = each["category"] + if category not in category_record: + category_record[category] = {"corr": 0.0, "wrong": 0.0} + category_record["random"] = {"corr": 0.0, "wrong": 0.0} + if not each["pred"]: + random.seed(12345) + x = random.randint(0, len(each["options"]) - 1) + if x == each["answer_index"]: + category_record[category]["corr"] += 1 + category_record["random"]["corr"] += 1 + else: + category_record[category]["wrong"] += 1 + category_record["random"]["wrong"] += 1 + elif each["pred"] == each["answer"]: + category_record[category]["corr"] += 1 + else: + category_record[category]["wrong"] += 1 + success = True + except Exception as e: + print("Error", e) + return res, category_record + + +def evaluate(subjects): + test_df, dev_df = load_mmlu_pro() + if not subjects: + subjects = list(test_df.keys()) + print("assigned subjects", subjects) + lock = threading.Lock() + system_prompt = config["inference"]["system_prompt"] + for subject in subjects: + start = time.time() + print(f"Testing {subject}...") + config["inference"]["system_prompt"] = system_prompt.replace( + "{subject}", subject + ) + test_data = test_df[subject] + output_res_path = os.path.join(output_dir, subject + "_result.json") + output_summary_path = os.path.join(output_dir, subject + "_summary.json") + res, category_record = update_result(output_res_path, lock) + + with ThreadPoolExecutor(max_workers=config["test"]["parallel"]) as executor: + futures = { + executor.submit(run_single_question, each, dev_df, res): each + for each in test_data + } + for future in tqdm( + as_completed(futures), total=len(futures), smoothing=0.0, ascii=True + ): + each = futures[future] + label = each["answer"] + category = subject + prompt, response, pred, exist = future.result() + if exist: + continue + if response is not None: + res, category_record = update_result(output_res_path, lock) + if category not in category_record: + category_record[category] = {"corr": 0.0, "wrong": 0.0} + if config["log"]["log_prompt"]: + each["prompt"] = prompt + each["response"] = response + each["pred"] = pred + res.append(each) + if config["log"]["verbosity"] >= 2: + log_json = { + "id": each["question_id"], + "question": each["question"], + "response": each["response"], + "pred": each["pred"], + "answer": each["answer"], + } + print("\n" + json.dumps(log_json, indent="\t")) + if pred is not None: + if pred == label: + category_record[category]["corr"] += 1 + else: + category_record[category]["wrong"] += 1 + else: + category_record[category]["wrong"] += 1 + save_res(res, output_res_path, lock) + save_summary(category_record, output_summary_path, lock) + res, category_record = update_result(output_res_path, lock) + save_res(res, output_res_path, lock) + hours, minutes, seconds = elapsed(start) + log( + f"Finished testing {subject} in {hours} hours, {minutes} minutes, {seconds} seconds." + ) + save_summary(category_record, output_summary_path, lock, report=True) + + +def save_res(res, output_res_path, lock): + temp = [] + exist_q_id = [] + for each in res: + if each["question_id"] not in exist_q_id: + exist_q_id.append(each["question_id"]) + temp.append(each) + else: + continue + res = temp + with lock: + with open(output_res_path, "w") as fo: + fo.write(json.dumps(res, indent="\t")) + + +def print_score(label, corr, wrong): + try: + corr = int(corr) + wrong = int(wrong) + total = corr + wrong + acc = corr / total * 100 + log(f"{label}, {corr}/{total}, {acc:.2f}%") + except Exception as e: + log(f"{label}, {e} error") + + +def save_summary(category_record, output_summary_path, lock, report=False): + total_corr = 0.0 + total_wrong = 0.0 + for k, v in category_record.items(): + if k == "total" or k == "random": + continue + cat_acc = v["corr"] / (v["corr"] + v["wrong"]) + category_record[k]["acc"] = cat_acc + total_corr += v["corr"] + total_wrong += v["wrong"] + acc = total_corr / (total_corr + total_wrong) + category_record["total"] = {"corr": total_corr, "wrong": total_wrong, "acc": acc} + if report: + print_score("Total", total_corr, total_wrong) + if "random" in category_record: + random_corr = category_record["random"]["corr"] + random_wrong = category_record["random"]["wrong"] + print_score( + "Random Guess Attempts", + random_corr + random_wrong, + total_corr + total_wrong - random_corr - random_wrong, + ) + print_score("Correct Random Guesses", random_corr, random_wrong) + print_score( + "Adjusted Score Without Random Guesses", + total_corr - random_corr, + total_wrong - random_wrong, + ) + with lock: + with open(output_summary_path, "w") as fo: + fo.write(json.dumps(category_record, indent="\t")) + + +def final_report(assigned_subjects): + total_corr = 0.0 + total_wrong = 0.0 + random_corr = 0.0 + random_wrong = 0.0 + names = ["overall"] + assigned_subjects + table = "| " + " | ".join(names) + " |\n" + separators = [re.sub(r".", "-", name) for name in names] + table += "| " + " | ".join(separators) + " |\n" + scores = [] + for file in assigned_subjects: + res = json.load(open(os.path.join(output_dir, file + "_summary.json"))) + cat_corr = res["total"]["corr"] + total_corr += cat_corr + cat_wrong = res["total"]["wrong"] + total_wrong += cat_wrong + scores.append(cat_corr / (cat_corr + cat_wrong)) + if "random" in res: + random_corr += res["random"]["corr"] + random_wrong += res["random"]["wrong"] + print_score("Total", total_corr, total_wrong) + if random_corr and random_wrong: + print_score( + "Random Guess Attempts", + random_corr + random_wrong, + total_corr + total_wrong - random_corr - random_wrong, + ) + print_score("Correct Random Guesses", random_corr, random_wrong) + print_score( + "Adjusted Score Without Random Guesses", + total_corr - random_corr, + total_wrong - random_wrong, + ) + scores.insert(0, total_corr / (total_corr + total_wrong)) + scores = [f"{score*100:.2f}" for score in scores] + table += "| " + " | ".join(scores) + " |" + token_report() + log("Markdown Table:") + log(table) + + +def elapsed(start): + duration = time.time() - start + duration_td = timedelta(seconds=duration) + hours, remainder = divmod(duration_td.seconds, 3600) + minutes, seconds = divmod(remainder, 60) + return hours, minutes, seconds + + +def token_report(): + ptoks = [] + ctoks = [] + while not usage_q.empty(): + usage = usage_q.get() + ptoks.append(usage[0]) + ctoks.append(usage[1]) + if ptoks and ctoks: + log("Token Usage:") + duration = end - start + ptoks = np.array(ptoks) + ctoks = np.array(ctoks) + log( + f"Prompt tokens: min {ptoks.min()}, average {ptoks.mean():.0f}, max {ptoks.max()}, total {ptoks.sum()}, tk/s {ptoks.sum()/duration:.2f}" + ) + log( + f"Completion tokens: min {ctoks.min()}, average {ctoks.mean():.0f}, max {ctoks.max()}, total {ctoks.sum()}, tk/s {ctoks.sum()/duration:.2f}" + ) + + +if __name__ == "__main__": + usage_q = queue.Queue() + output_dir = "eval_results/" + re.sub(r"\W", "-", config["server"]["model"]) + os.makedirs(output_dir, exist_ok=True) + log_path = os.path.join(output_dir, "report.txt") + try: + os.remove(log_path) + except: + pass + config_copy = copy.deepcopy(config) + del config_copy["server"]["api_key"] + del config_copy["test"]["categories"] + log(f"{datetime.now()}") + log(json.dumps(config_copy, indent="\t")) + assigned_subjects = config["test"]["categories"] + start = time.time() + evaluate(assigned_subjects) + end = time.time() + hours, minutes, seconds = elapsed(start) + log( + f"Finished the benchmark in {hours} hours, {minutes} minutes, {seconds} seconds." + ) + final_report(assigned_subjects) + print("Report saved to:", log_path) diff --git a/App_Function_Libraries/Benchmarks_Evaluations/__init__.py b/App_Function_Libraries/Benchmarks_Evaluations/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/App_Function_Libraries/Benchmarks_Evaluations/ms_g_eval.py b/App_Function_Libraries/Benchmarks_Evaluations/ms_g_eval.py new file mode 100644 index 0000000000000000000000000000000000000000..9cbcfacd1e0ec6005fe7aa231cadb28671ab6cad --- /dev/null +++ b/App_Function_Libraries/Benchmarks_Evaluations/ms_g_eval.py @@ -0,0 +1,498 @@ +####################################################################################################################### +# +# Evaluations_Benchmarks_tab.py +# +# Description: This file contains the code to evaluate the generated text using G-Eval metric. +# +# Scripts taken from https://github.com/microsoft/promptflow/tree/main/examples/flows/evaluation/eval-summarization and modified. +# +import configparser +import inspect +import json +import logging +import os +import re +from typing import Dict, Callable, List, Any + +import gradio as gr +from tenacity import ( + RetryError, + Retrying, + after_log, + before_sleep_log, + stop_after_attempt, + wait_random_exponential, +) + +from App_Function_Libraries.Chat import chat_api_call + +# +####################################################################################################################### +# +# Start of G-Eval.py + +logger = logging.getLogger(__name__) + +current_dir = os.path.dirname(os.path.abspath(__file__)) +# Construct the path to the config file +config_path = os.path.join(current_dir, 'Config_Files', 'config.txt') +# Read the config file +config = configparser.ConfigParser() +config.read(config_path) + + +def aggregate( + fluency_list: List[float], + consistency_list: List[float], + relevance_list: List[float], + coherence_list: List[float], +) -> Dict[str, float]: + """ + Takes list of scores for 4 dims and outputs average for them. + + Args: + fluency_list (List(float)): list of fluency scores + consistency_list (List(float)): list of consistency scores + relevance_list (List(float)): list of relevance scores + coherence_list (List(float)): list of coherence scores + + Returns: + Dict[str, float]: Returns average scores + """ + average_fluency = sum(fluency_list) / len(fluency_list) + average_consistency = sum(consistency_list) / len(consistency_list) + average_relevance = sum(relevance_list) / len(relevance_list) + average_coherence = sum(coherence_list) / len(coherence_list) + + log_metric("average_fluency", average_fluency) + log_metric("average_consistency", average_consistency) + log_metric("average_relevance", average_relevance) + log_metric("average_coherence", average_coherence) + + return { + "average_fluency": average_fluency, + "average_consistency": average_consistency, + "average_relevance": average_relevance, + "average_coherence": average_coherence, + } + +def run_geval(transcript: str, summary: str, api_key: str, api_name: str = None, save: bool = False): + try: + validate_inputs(transcript, summary, api_name, api_key) + except ValueError as e: + return str(e) + + prompts = { + "coherence": """You will be given one summary written for a source document. + + Your task is to rate the summary on one metric. + + Please make sure you read and understand these instructions carefully. Please keep this document open while reviewing, and refer to it as needed. + + Evaluation Criteria: + + Coherence (1-5) - the collective quality of all sentences. We align this dimension with the DUC quality question of structure and coherence whereby "the summary should be well-structured and well-organized. The summary should not just be a heap of related information, but should build from sentence to a coherent body of information about a topic." + + Evaluation Steps: + + 1. Read the source document carefully and identify the main topic and key points. + 2. Read the summary and compare it to the source document. Check if the summary covers the main topic and key points of the source document, and if it presents them in a clear and logical order. + 3. Assign a score for coherence on a scale of 1 to 5, where 1 is the lowest and 5 is the highest based on the Evaluation Criteria. + + + Example: + + + Source Document: + + {{Document}} + + Summary: + + {{Summary}} + + + Evaluation Form (scores ONLY): + + - Coherence:""", + "consistency": """You will be given a source document. You will then be given one summary written for this source document. + + Your task is to rate the summary on one metric. + + Please make sure you read and understand these instructions carefully. Please keep this document open while reviewing, and refer to it as needed. + + + Evaluation Criteria: + + Consistency (1-5) - the factual alignment between the summary and the summarized source. A factually consistent summary contains only statements that are entailed by the source document. Annotators were also asked to penalize summaries that contained hallucinated facts. + + Evaluation Steps: + + 1. Read the source document carefully and identify the main facts and details it presents. + 2. Read the summary and compare it to the source document. Check if the summary contains any factual errors that are not supported by the source document. + 3. Assign a score for consistency based on the Evaluation Criteria. + + + Example: + + + Source Document: + + {{Document}} + + Summary: + + {{Summary}} + + + Evaluation Form (scores ONLY): + + - Consistency:""", + "fluency": """You will be given one summary written for a source document. + + Your task is to rate the summary on one metric. + + Please make sure you read and understand these instructions carefully. Please keep this document open while reviewing, and refer to it as needed. + + + Evaluation Criteria: + + Fluency (1-3): the quality of the summary in terms of grammar, spelling, punctuation, word choice, and sentence structure. + + - 1: Poor. The summary has many errors that make it hard to understand or sound unnatural. + - 2: Fair. The summary has some errors that affect the clarity or smoothness of the text, but the main points are still comprehensible. + - 3: Good. The summary has few or no errors and is easy to read and follow. + + + Example: + + Summary: + + {{Summary}} + + + Evaluation Form (scores ONLY): + + - Fluency (1-3):""", + "relevance": """You will be given one summary written for a source document. + + Your task is to rate the summary on one metric. + + Please make sure you read and understand these instructions carefully. Please keep this document open while reviewing, and refer to it as needed. + + Evaluation Criteria: + + Relevance (1-5) - selection of important content from the source. The summary should include only important information from the source document. Annotators were instructed to penalize summaries which contained redundancies and excess information. + + Evaluation Steps: + + 1. Read the summary and the source document carefully. + 2. Compare the summary to the source document and identify the main points of the source document. + 3. Assess how well the summary covers the main points of the source document, and how much irrelevant or redundant information it contains. + 4. Assign a relevance score from 1 to 5. + + + Example: + + + Source Document: + + {{Document}} + + Summary: + + {{Summary}} + + + Evaluation Form (scores ONLY): + + - Relevance:""" + } + + scores = {} + explanations = {} + for metric, prompt in prompts.items(): + full_prompt = prompt.replace("{{Document}}", transcript).replace("{{Summary}}", summary) + try: + score = geval_summarization(full_prompt, 5 if metric != "fluency" else 3, api_name, api_key) + scores[metric] = score + explanations[metric] = "Score based on the evaluation criteria." + except Exception as e: + error_message = detailed_api_error(api_name, e) + return error_message + + avg_scores = aggregate([scores['fluency']], [scores['consistency']], + [scores['relevance']], [scores['coherence']]) + + results = { + "scores": scores, + "average_scores": avg_scores + } + logging.debug("Results: %s", results) + + if save is not None: + logging.debug("Saving results to geval_results.json") + save_eval_results(results) + logging.debug("Results saved to geval_results.json") + + formatted_result = f""" + Confabulation Check Results: + + Coherence: {scores['coherence']:.2f} - {explanations['coherence']} + Consistency: {scores['consistency']:.2f} - {explanations['consistency']} + Fluency: {scores['fluency']:.2f} - {explanations['fluency']} + Relevance: {scores['relevance']:.2f} - {explanations['relevance']} + + Overall Assessment: The summary has been evaluated on four key metrics. + The average scores are: + Fluency: {avg_scores['average_fluency']:.2f} + Consistency: {avg_scores['average_consistency']:.2f} + Relevance: {avg_scores['average_relevance']:.2f} + Coherence: {avg_scores['average_coherence']:.2f} + + These scores indicate the overall quality of the summary in terms of its + coherence, consistency with the original text, fluency of language, and + relevance of content. + """ + + return formatted_result + + +def create_geval_tab(): + with gr.Tab("G-Eval", id="g-eval"): + gr.Markdown("# G-Eval Summarization Evaluation") + with gr.Row(): + with gr.Column(): + document_input = gr.Textbox(label="Source Document", lines=10) + summary_input = gr.Textbox(label="Summary", lines=5) + api_name_input = gr.Dropdown( + choices=["OpenAI", "Anthropic", "Cohere", "Groq", "OpenRouter", "DeepSeek", "HuggingFace", "Mistral", "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM", "Local-LLM", "Ollama"], + label="Select API" + ) + api_key_input = gr.Textbox(label="API Key (if required)", type="password") + save_value = gr.Checkbox(label="Save Results to a JSON file(geval_results.json)") + evaluate_button = gr.Button("Evaluate Summary") + with gr.Column(): + output = gr.Textbox(label="Evaluation Results", lines=10) + + evaluate_button.click( + fn=run_geval, + inputs=[document_input, summary_input, api_name_input, api_key_input, save_value], + outputs=output + ) + + return document_input, summary_input, api_name_input, api_key_input, evaluate_button, output + + +def parse_output(output: str, max: float) -> float: + """ + Function that extracts numerical score from the beginning of string + + Args: + output (str): String to search + max (float): Maximum score allowed + + Returns: + float: The extracted score + """ + matched: List[str] = re.findall(r"(? max: + raise ValueError(f"Parsed number: {score} was larger than max score: {max}") + else: + raise ValueError(f"More than one number detected in input. Input to parser was: {output}") + else: + raise ValueError(f'No number detected in input. Input to parser was "{output}". ') + return score + +def geval_summarization( + prompt_with_src_and_gen: str, + max_score: float, + api_endpoint: str, + api_key: str, +) -> float: + model = get_model_from_config(api_endpoint) + + try: + for attempt in Retrying( + reraise=True, + before_sleep=before_sleep_log(logger, logging.INFO), + after=after_log(logger, logging.INFO), + wait=wait_random_exponential(multiplier=1, min=1, max=120), + stop=stop_after_attempt(10), + ): + with attempt: + system_message="You are a helpful AI assistant" + # TEMP setting for Confabulation check + temp = 0.7 + logging.info(f"Debug - geval_summarization Function - API Endpoint: {api_endpoint}") + try: + response = chat_api_call(api_endpoint, api_key, prompt_with_src_and_gen, "", temp, system_message) + except Exception as e: + raise ValueError(f"Unsupported API endpoint: {api_endpoint}") + except RetryError: + logger.exception(f"geval {api_endpoint} call failed\nInput prompt was: {prompt_with_src_and_gen}") + raise + + try: + score = parse_output(response, max_score) + except ValueError as e: + logger.warning(f"Error parsing output: {e}") + score = 0 + + return score + + +def get_model_from_config(api_name: str) -> str: + model = config.get('models', api_name) + if isinstance(model, dict): + # If the model is a dictionary, return a specific key or a default value + return model.get('name', str(model)) # Adjust 'name' to the appropriate key if needed + return str(model) if model is not None else "" + +def aggregate_llm_scores(llm_responses: List[str], max_score: float) -> float: + """Parse and average valid scores from the generated responses of + the G-Eval LLM call. + + Args: + llm_responses (List[str]): List of scores from multiple LLMs + max_score (float): The maximum score allowed. + + Returns: + float: The average of all the valid scores + """ + all_scores = [] + error_count = 0 + for generated in llm_responses: + try: + parsed = parse_output(generated, max_score) + all_scores.append(parsed) + except ValueError as e: + logger.warning(e) + error_count += 1 + if error_count: + logger.warning(f"{error_count} out of 20 scores were discarded due to corrupt g-eval generation") + score = sum(all_scores) / len(all_scores) + return score + + +def validate_inputs(document: str, summary: str, api_name: str, api_key: str) -> None: + """ + Validate inputs for the G-Eval function. + + Args: + document (str): The source document + summary (str): The summary to evaluate + api_name (str): The name of the API to use + api_key (str): The API key + + Raises: + ValueError: If any of the inputs are invalid + """ + if not document.strip(): + raise ValueError("Source document cannot be empty") + if not summary.strip(): + raise ValueError("Summary cannot be empty") + if api_name.lower() not in ["openai", "anthropic", "cohere", "groq", "openrouter", "deepseek", "huggingface", + "mistral", "llama.cpp", "kobold", "ooba", "tabbyapi", "vllm", "local-llm", "ollama"]: + raise ValueError(f"Unsupported API: {api_name}") + + +def detailed_api_error(api_name: str, error: Exception) -> str: + """ + Generate a detailed error message for API failures. + + Args: + api_name (str): The name of the API that failed + error (Exception): The exception that was raised + + Returns: + str: A detailed error message + """ + error_type = type(error).__name__ + error_message = str(error) + return f"API Failure: {api_name}\nError Type: {error_type}\nError Message: {error_message}\nPlease check your API key and network connection, and try again." + + +def save_eval_results(results: Dict[str, Any], filename: str = "geval_results.json") -> None: + """ + Save evaluation results to a JSON file. + + Args: + results (Dict[str, Any]): The evaluation results + filename (str): The name of the file to save results to + """ + with open(filename, 'w') as f: + json.dump(results, f, indent=2) + print(f"Results saved to {filename}") + + + + +# +# +####################################################################################################################### +# +# Taken from: https://github.com/microsoft/promptflow/blob/b5a68f45e4c3818a29e2f79a76f2e73b8ea6be44/src/promptflow-core/promptflow/_core/metric_logger.py + +class MetricLoggerManager: + _instance = None + + def __init__(self): + self._metric_loggers = [] + + @staticmethod + def get_instance() -> "MetricLoggerManager": + if MetricLoggerManager._instance is None: + MetricLoggerManager._instance = MetricLoggerManager() + return MetricLoggerManager._instance + + def log_metric(self, key, value, variant_id=None): + for logger in self._metric_loggers: + if len(inspect.signature(logger).parameters) == 2: + logger(key, value) # If the logger only accepts two parameters, we don't pass variant_id + else: + logger(key, value, variant_id) + + def add_metric_logger(self, logger_func: Callable): + existing_logger = next((logger for logger in self._metric_loggers if logger is logger_func), None) + if existing_logger: + return + if not callable(logger_func): + return + sign = inspect.signature(logger_func) + # We accept two kinds of metric loggers: + # def log_metric(k, v) + # def log_metric(k, v, variant_id) + if len(sign.parameters) not in [2, 3]: + return + self._metric_loggers.append(logger_func) + + def remove_metric_logger(self, logger_func: Callable): + self._metric_loggers.remove(logger_func) + + +def log_metric(key, value, variant_id=None): + """Log a metric for current promptflow run. + + :param key: Metric name. + :type key: str + :param value: Metric value. + :type value: float + :param variant_id: Variant id for the metric. + :type variant_id: str + """ + MetricLoggerManager.get_instance().log_metric(key, value, variant_id) + + +def add_metric_logger(logger_func: Callable): + MetricLoggerManager.get_instance().add_metric_logger(logger_func) + + +def remove_metric_logger(logger_func: Callable): + MetricLoggerManager.get_instance().remove_metric_logger(logger_func) +# +# End of G-Eval.py +####################################################################################################################### \ No newline at end of file diff --git a/App_Function_Libraries/Books/.pytest_cache/.gitignore b/App_Function_Libraries/Books/.pytest_cache/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..08a7f458f1f002823bc794c47ca1996a57e72c86 --- /dev/null +++ b/App_Function_Libraries/Books/.pytest_cache/.gitignore @@ -0,0 +1,2 @@ +# Created by pytest automatically. +* diff --git a/App_Function_Libraries/Books/.pytest_cache/CACHEDIR.TAG b/App_Function_Libraries/Books/.pytest_cache/CACHEDIR.TAG new file mode 100644 index 0000000000000000000000000000000000000000..fce15ad7eaa74e5682b644c84efb75334c112f95 --- /dev/null +++ b/App_Function_Libraries/Books/.pytest_cache/CACHEDIR.TAG @@ -0,0 +1,4 @@ +Signature: 8a477f597d28d172789f06886806bc55 +# This file is a cache directory tag created by pytest. +# For information about cache directory tags, see: +# https://bford.info/cachedir/spec.html diff --git a/App_Function_Libraries/Books/.pytest_cache/README.md b/App_Function_Libraries/Books/.pytest_cache/README.md new file mode 100644 index 0000000000000000000000000000000000000000..c7526af2448672de4537dfed042ed74daadb17bf --- /dev/null +++ b/App_Function_Libraries/Books/.pytest_cache/README.md @@ -0,0 +1,8 @@ +# pytest cache directory # + +This directory contains data from the pytest's cache plugin, +which provides the `--lf` and `--ff` options, as well as the `cache` fixture. + +**Do not** commit this to version control. + +See [the docs](https://docs.pytest.org/en/stable/how-to/cache.html) for more information. diff --git a/App_Function_Libraries/Books/.pytest_cache/v/cache/lastfailed b/App_Function_Libraries/Books/.pytest_cache/v/cache/lastfailed new file mode 100644 index 0000000000000000000000000000000000000000..b092cbef463fb2b33edf5f000b4faf34cb568129 --- /dev/null +++ b/App_Function_Libraries/Books/.pytest_cache/v/cache/lastfailed @@ -0,0 +1,10 @@ +{ + "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_file": true, + "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_missing_metadata": true, + "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_with_auto_summarize": true, + "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_process_zip_file": true, + "test_Book_Ingestion_tab.py::TestBookIngestionTab::test_import_epub_file": true, + "test_Book_Ingestion_tab.py::TestBookIngestionTab::test_import_zip_file": true, + "test_Book_Ingestion_lib.py": true, + "test_Book_Ingestion_tab.py": true +} \ No newline at end of file diff --git a/App_Function_Libraries/Books/.pytest_cache/v/cache/nodeids b/App_Function_Libraries/Books/.pytest_cache/v/cache/nodeids new file mode 100644 index 0000000000000000000000000000000000000000..a14807f56324ae1672807b67ee5a676d1fe0c1d2 --- /dev/null +++ b/App_Function_Libraries/Books/.pytest_cache/v/cache/nodeids @@ -0,0 +1,11 @@ +[ + "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_file", + "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_invalid_file", + "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_missing_metadata", + "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_import_epub_with_auto_summarize", + "test_Book_Ingestion_lib.py::TestBookIngestionTab::test_process_zip_file", + "test_Book_Ingestion_tab.py::TestBookIngestionTab::test_import_epub_file", + "test_Book_Ingestion_tab.py::TestBookIngestionTab::test_import_no_file", + "test_Book_Ingestion_tab.py::TestBookIngestionTab::test_import_unsupported_file", + "test_Book_Ingestion_tab.py::TestBookIngestionTab::test_import_zip_file" +] \ No newline at end of file diff --git a/App_Function_Libraries/Books/.pytest_cache/v/cache/stepwise b/App_Function_Libraries/Books/.pytest_cache/v/cache/stepwise new file mode 100644 index 0000000000000000000000000000000000000000..0637a088a01e8ddab3bf3fa98dbe804cbde1a0dc --- /dev/null +++ b/App_Function_Libraries/Books/.pytest_cache/v/cache/stepwise @@ -0,0 +1 @@ +[] \ No newline at end of file diff --git a/App_Function_Libraries/Books/Book_Ingestion_Lib.py b/App_Function_Libraries/Books/Book_Ingestion_Lib.py new file mode 100644 index 0000000000000000000000000000000000000000..66e49d904c65c2839a31d06edf5054d96c8c7fb6 --- /dev/null +++ b/App_Function_Libraries/Books/Book_Ingestion_Lib.py @@ -0,0 +1,577 @@ +# Book_Ingestion_Lib.py +######################################### +# Library to hold functions for ingesting book files.# +# +#################### +# Function List +# +# 1. ingest_text_file(file_path, title=None, author=None, keywords=None): +# 2. +# +# +#################### +# +# Imports +import os +import re +import tempfile +import zipfile +from datetime import datetime +import logging +# +# External Imports +import ebooklib +from bs4 import BeautifulSoup +from ebooklib import epub +# +# Import Local +from App_Function_Libraries.DB.DB_Manager import add_media_with_keywords, add_media_to_database +from App_Function_Libraries.Summarization.Summarization_General_Lib import perform_summarization +from App_Function_Libraries.Chunk_Lib import chunk_ebook_by_chapters +from App_Function_Libraries.Metrics.metrics_logger import log_counter, log_histogram +# +####################################################################################################################### +# Function Definitions +# + +def import_epub(file_path, + title=None, + author=None, + keywords=None, + custom_prompt=None, + system_prompt=None, + summary=None, + auto_summarize=False, + api_name=None, + api_key=None, + chunk_options=None, + custom_chapter_pattern=None + ): + """ + Imports an EPUB file, extracts its content, chunks it, optionally summarizes it, and adds it to the database. + + Parameters: + - file_path (str): Path to the EPUB file. + - title (str, optional): Title of the book. + - author (str, optional): Author of the book. + - keywords (str, optional): Comma-separated keywords for the book. + - custom_prompt (str, optional): Custom user prompt for summarization. + - summary (str, optional): Predefined summary of the book. + - auto_summarize (bool, optional): Whether to auto-summarize the chunks. + - api_name (str, optional): API name for summarization. + - api_key (str, optional): API key for summarization. + - chunk_options (dict, optional): Options for chunking. + - custom_chapter_pattern (str, optional): Custom regex pattern for chapter detection. + + Returns: + - str: Status message indicating success or failure. + """ + try: + logging.info(f"Importing EPUB file from {file_path}") + log_counter("epub_import_attempt", labels={"file_path": file_path}) + + start_time = datetime.now() + + # Convert EPUB to Markdown + markdown_content = epub_to_markdown(file_path) + logging.debug("Converted EPUB to Markdown.") + + # Extract metadata if not provided + if not title or not author: + extracted_title, extracted_author = extract_epub_metadata(markdown_content) + title = title or extracted_title or os.path.splitext(os.path.basename(file_path))[0] + author = author or extracted_author or "Unknown" + logging.debug(f"Extracted metadata - Title: {title}, Author: {author}") + + # Process keywords + keyword_list = [kw.strip() for kw in keywords.split(',')] if keywords else [] + logging.debug(f"Keywords: {keyword_list}") + + # Set default chunk options if not provided + if chunk_options is None: + chunk_options = { + 'method': 'chapter', + 'max_size': 500, + 'overlap': 200, + 'custom_chapter_pattern': custom_chapter_pattern + } + else: + # Ensure 'method' is set to 'chapter' when using chapter chunking + chunk_options.setdefault('method', 'chapter') + chunk_options.setdefault('custom_chapter_pattern', custom_chapter_pattern) + + # Chunk the content by chapters + chunks = chunk_ebook_by_chapters(markdown_content, chunk_options) + logging.info(f"Total chunks created: {len(chunks)}") + log_histogram("epub_chunks_created", len(chunks), labels={"file_path": file_path}) + + if chunks: + logging.debug(f"Structure of first chunk: {chunks[0].keys()}") + + # Handle summarization if enabled + if auto_summarize and api_name and api_key: + logging.info("Auto-summarization is enabled.") + summarized_chunks = [] + for chunk in chunks: + chunk_text = chunk.get('text', '') + if chunk_text: + summary_text = perform_summarization(api_name, chunk_text, custom_prompt, api_key, + recursive_summarization=False, temp=None, + system_message=system_prompt + ) + chunk['metadata']['summary'] = summary_text + summarized_chunks.append(chunk) + chunks = summarized_chunks + logging.info("Summarization of chunks completed.") + log_counter("epub_chunks_summarized", value=len(chunks), labels={"file_path": file_path}) + else: + # If not summarizing, set a default summary or use provided summary + if summary: + logging.debug("Using provided summary.") + else: + summary = "No summary provided." + + # Create info_dict + info_dict = { + 'title': title, + 'uploader': author, + 'ingestion_date': datetime.now().strftime('%Y-%m-%d') + } + + # Prepare segments for database + segments = [{'Text': chunk.get('text', chunk.get('content', ''))} for chunk in chunks] + logging.debug(f"Prepared segments for database. Number of segments: {len(segments)}") + + # Add to database + result = add_media_to_database( + url=file_path, + info_dict=info_dict, + segments=segments, + summary=summary, + keywords=keyword_list, + custom_prompt_input=custom_prompt, + whisper_model="Imported", + media_type="ebook", + overwrite=False + ) + + end_time = datetime.now() + processing_time = (end_time - start_time).total_seconds() + log_histogram("epub_import_duration", processing_time, labels={"file_path": file_path}) + + logging.info(f"Ebook '{title}' by {author} imported successfully. Database result: {result}") + log_counter("epub ingested into the DB successfully", labels={"file_path": file_path}) + return f"Ebook '{title}' by {author} imported successfully. Database result: {result}" + + except Exception as e: + logging.exception(f"Error importing ebook: {str(e)}") + log_counter("epub_import_error", labels={"file_path": file_path, "error": str(e)}) + return f"Error importing ebook: {str(e)}" + + +# FIXME +def process_zip_file(zip_file, + title, + author, + keywords, + custom_prompt, + system_prompt, + summary, + auto_summarize, + api_name, + api_key, + chunk_options + ): + """ + Processes a ZIP file containing multiple EPUB files and imports each one. + + Parameters: + - zip_file (file-like object): The ZIP file to process. + - title (str): Title prefix for the books. + - author (str): Author name for the books. + - keywords (str): Comma-separated keywords. + - custom_prompt (str): Custom user prompt for summarization. + - summary (str): Predefined summary (not used in this context). + - auto_summarize (bool): Whether to auto-summarize the chunks. + - api_name (str): API name for summarization. + - api_key (str): API key for summarization. + - chunk_options (dict): Options for chunking. + + Returns: + - str: Combined status messages for all EPUB files in the ZIP. + """ + results = [] + try: + with tempfile.TemporaryDirectory() as temp_dir: + zip_path = zip_file.name if hasattr(zip_file, 'name') else zip_file.path + logging.info(f"Extracting ZIP file {zip_path} to temporary directory {temp_dir}") + log_counter("zip_processing_attempt", labels={"zip_path": zip_path}) + + with zipfile.ZipFile(zip_path, 'r') as zip_ref: + zip_ref.extractall(temp_dir) + + epub_files = [f for f in os.listdir(temp_dir) if f.lower().endswith('.epub')] + log_histogram("epub_files_in_zip", len(epub_files), labels={"zip_path": zip_path}) + + for filename in epub_files: + file_path = os.path.join(temp_dir, filename) + logging.info(f"Processing EPUB file {filename} from ZIP.") + result = import_epub( + file_path=file_path, + title=title, + author=author, + keywords=keywords, + custom_prompt=custom_prompt, + summary=summary, + auto_summarize=auto_summarize, + api_name=api_name, + api_key=api_key, + chunk_options=chunk_options, + custom_chapter_pattern=chunk_options.get('custom_chapter_pattern') if chunk_options else None + ) + results.append(f"File: {filename} - {result}") + + logging.info("Completed processing all EPUB files in the ZIP.") + log_counter("zip_processing_success", labels={"zip_path": zip_path}) + except Exception as e: + logging.exception(f"Error processing ZIP file: {str(e)}") + log_counter("zip_processing_error", labels={"zip_path": zip_path, "error": str(e)}) + return f"Error processing ZIP file: {str(e)}" + + return "\n".join(results) + + +def import_file_handler(file, + title, + author, + keywords, + system_prompt, + custom_prompt, + auto_summarize, + api_name, + api_key, + max_chunk_size, + chunk_overlap, + custom_chapter_pattern + ): + try: + log_counter("file_import_attempt", labels={"file_name": file.name}) + + # Handle max_chunk_size + if isinstance(max_chunk_size, str): + max_chunk_size = int(max_chunk_size) if max_chunk_size.strip() else 4000 + elif not isinstance(max_chunk_size, int): + max_chunk_size = 4000 # Default value if not a string or int + + # Handle chunk_overlap + if isinstance(chunk_overlap, str): + chunk_overlap = int(chunk_overlap) if chunk_overlap.strip() else 0 + elif not isinstance(chunk_overlap, int): + chunk_overlap = 0 # Default value if not a string or int + + chunk_options = { + 'method': 'chapter', + 'max_size': max_chunk_size, + 'overlap': chunk_overlap, + 'custom_chapter_pattern': custom_chapter_pattern if custom_chapter_pattern else None + } + + if file is None: + log_counter("file_import_error", labels={"error": "No file uploaded"}) + return "No file uploaded." + + file_path = file.name + if not os.path.exists(file_path): + log_counter("file_import_error", labels={"error": "File not found", "file_name": file.name}) + return "Uploaded file not found." + + start_time = datetime.now() + + if file_path.lower().endswith('.epub'): + status = import_epub( + file_path, + title, + author, + keywords, + custom_prompt=custom_prompt, + system_prompt=system_prompt, + summary=None, + auto_summarize=auto_summarize, + api_name=api_name, + api_key=api_key, + chunk_options=chunk_options, + custom_chapter_pattern=custom_chapter_pattern + ) + log_counter("epub_import_success", labels={"file_name": file.name}) + result = f"📚 EPUB Imported Successfully:\n{status}" + elif file.name.lower().endswith('.zip'): + status = process_zip_file( + zip_file=file, + title=title, + author=author, + keywords=keywords, + custom_prompt=custom_prompt, + system_prompt=system_prompt, + summary=None, + auto_summarize=auto_summarize, + api_name=api_name, + api_key=api_key, + chunk_options=chunk_options + ) + log_counter("zip_import_success", labels={"file_name": file.name}) + result = f"📦 ZIP Processed Successfully:\n{status}" + elif file.name.lower().endswith(('.chm', '.html', '.pdf', '.xml', '.opml')): + file_type = file.name.split('.')[-1].upper() + log_counter("unsupported_file_type", labels={"file_type": file_type}) + result = f"{file_type} file import is not yet supported." + else: + log_counter("unsupported_file_type", labels={"file_type": file.name.split('.')[-1]}) + result = "❌ Unsupported file type. Please upload an `.epub` file or a `.zip` file containing `.epub` files." + + end_time = datetime.now() + processing_time = (end_time - start_time).total_seconds() + log_histogram("file_import_duration", processing_time, labels={"file_name": file.name}) + + return result + + except ValueError as ve: + logging.exception(f"Error parsing input values: {str(ve)}") + log_counter("file_import_error", labels={"error": "Invalid input", "file_name": file.name}) + return f"❌ Error: Invalid input for chunk size or overlap. Please enter valid numbers." + except Exception as e: + logging.exception(f"Error during file import: {str(e)}") + log_counter("file_import_error", labels={"error": str(e), "file_name": file.name}) + return f"❌ Error during import: {str(e)}" + + +def read_epub(file_path): + """ + Reads and extracts text from an EPUB file. + + Parameters: + - file_path (str): Path to the EPUB file. + + Returns: + - str: Extracted text content from the EPUB. + """ + try: + logging.info(f"Reading EPUB file from {file_path}") + book = epub.read_epub(file_path) + chapters = [] + for item in book.get_items(): + if item.get_type() == ebooklib.ITEM_DOCUMENT: + chapters.append(item.get_content()) + + text = "" + for html_content in chapters: + soup = BeautifulSoup(html_content, 'html.parser') + text += soup.get_text(separator='\n\n') + "\n\n" + logging.debug("EPUB content extraction completed.") + return text + except Exception as e: + logging.exception(f"Error reading EPUB file: {str(e)}") + raise + + +# Ingest a text file into the database with Title/Author/Keywords +def extract_epub_metadata(content): + title_match = re.search(r'Title:\s*(.*?)\n', content) + author_match = re.search(r'Author:\s*(.*?)\n', content) + + title = title_match.group(1) if title_match else None + author = author_match.group(1) if author_match else None + + return title, author + + +def ingest_text_file(file_path, title=None, author=None, keywords=None): + """ + Ingests a plain text file into the database with optional metadata. + + Parameters: + - file_path (str): Path to the text file. + - title (str, optional): Title of the document. + - author (str, optional): Author of the document. + - keywords (str, optional): Comma-separated keywords. + + Returns: + - str: Status message indicating success or failure. + """ + try: + with open(file_path, 'r', encoding='utf-8') as file: + content = file.read() + + # Check if it's a converted epub and extract metadata if so + if 'epub_converted' in (keywords or '').lower(): + extracted_title, extracted_author = extract_epub_metadata(content) + title = title or extracted_title + author = author or extracted_author + logging.debug(f"Extracted metadata for converted EPUB - Title: {title}, Author: {author}") + + # If title is still not provided, use the filename without extension + if not title: + title = os.path.splitext(os.path.basename(file_path))[0] + + # If author is still not provided, set it to 'Unknown' + if not author: + author = 'Unknown' + + # If keywords are not provided, use a default keyword + if not keywords: + keywords = 'text_file,epub_converted' + else: + keywords = f'text_file,epub_converted,{keywords}' + + # Add the text file to the database + add_media_with_keywords( + url=file_path, + title=title, + media_type='document', + content=content, + keywords=keywords, + prompt='No prompt for text files', + summary='No summary for text files', + transcription_model='None', + author=author, + ingestion_date=datetime.now().strftime('%Y-%m-%d') + ) + + logging.info(f"Text file '{title}' by {author} ingested successfully.") + return f"Text file '{title}' by {author} ingested successfully." + except Exception as e: + logging.error(f"Error ingesting text file: {str(e)}") + return f"Error ingesting text file: {str(e)}" + + +def ingest_folder(folder_path, keywords=None): + """ + Ingests all text files within a specified folder. + + Parameters: + - folder_path (str): Path to the folder containing text files. + - keywords (str, optional): Comma-separated keywords to add to each file. + + Returns: + - str: Combined status messages for all ingested text files. + """ + results = [] + try: + logging.info(f"Ingesting all text files from folder {folder_path}") + for filename in os.listdir(folder_path): + if filename.lower().endswith('.txt'): + file_path = os.path.join(folder_path, filename) + result = ingest_text_file(file_path, keywords=keywords) + results.append(result) + logging.info("Completed ingestion of all text files in the folder.") + except Exception as e: + logging.exception(f"Error ingesting folder: {str(e)}") + return f"Error ingesting folder: {str(e)}" + + return "\n".join(results) + + +def epub_to_markdown(epub_path): + """ + Converts an EPUB file to Markdown format, including the table of contents and chapter contents. + + Parameters: + - epub_path (str): Path to the EPUB file. + + Returns: + - str: Markdown-formatted content of the EPUB. + """ + try: + logging.info(f"Converting EPUB to Markdown from {epub_path}") + book = epub.read_epub(epub_path) + markdown_content = "# Table of Contents\n\n" + chapters = [] + + # Extract and format the table of contents + toc = book.toc + for item in toc: + if isinstance(item, tuple): + section, children = item + level = 1 + markdown_content += format_toc_item(section, level) + for child in children: + markdown_content += format_toc_item(child, level + 1) + else: + markdown_content += format_toc_item(item, 1) + + markdown_content += "\n---\n\n" + + # Process each chapter + for item in book.get_items(): + if item.get_type() == ebooklib.ITEM_DOCUMENT: + chapter_content = item.get_content().decode('utf-8') + soup = BeautifulSoup(chapter_content, 'html.parser') + + # Extract chapter title + title = soup.find(['h1', 'h2', 'h3']) + if title: + chapter_title = title.get_text() + markdown_content += f"# {chapter_title}\n\n" + + # Process chapter content + for elem in soup.find_all(['p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'ul', 'ol']): + if elem.name.startswith('h'): + level = int(elem.name[1]) + markdown_content += f"{'#' * level} {elem.get_text()}\n\n" + elif elem.name == 'p': + markdown_content += f"{elem.get_text()}\n\n" + elif elem.name in ['ul', 'ol']: + for li in elem.find_all('li'): + prefix = '-' if elem.name == 'ul' else '1.' + markdown_content += f"{prefix} {li.get_text()}\n" + markdown_content += "\n" + + markdown_content += "---\n\n" + + logging.debug("EPUB to Markdown conversion completed.") + return markdown_content + + except Exception as e: + logging.exception(f"Error converting EPUB to Markdown: {str(e)}") + raise + + +def format_toc_item(item, level): + """ + Formats a table of contents item into Markdown list format. + + Parameters: + - item (epub.Link or epub.Section): TOC item. + - level (int): Heading level for indentation. + + Returns: + - str: Markdown-formatted TOC item. + """ + try: + if isinstance(item, epub.Link): + title = item.title + elif isinstance(item, epub.Section): + title = item.title + else: + title = str(item) + + return f"{' ' * (level - 1)}- [{title}](#{slugify(title)})\n" + except Exception as e: + logging.exception(f"Error formatting TOC item: {str(e)}") + return "" + + +def slugify(text): + """ + Converts a string into a slug suitable for Markdown links. + + Parameters: + - text (str): The text to slugify. + + Returns: + - str: Slugified text. + """ + return re.sub(r'[\W_]+', '-', text.lower()).strip('-') + +# +# End of Function Definitions +####################################################################################################################### diff --git a/App_Function_Libraries/Books/__init__.py b/App_Function_Libraries/Books/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/App_Function_Libraries/Character_Chat/Character_Chat_Lib.py b/App_Function_Libraries/Character_Chat/Character_Chat_Lib.py new file mode 100644 index 0000000000000000000000000000000000000000..3dbc0e42e390e2182117483c7f99ca4ca5ae02f2 --- /dev/null +++ b/App_Function_Libraries/Character_Chat/Character_Chat_Lib.py @@ -0,0 +1,607 @@ +# Character_Chat_Lib.py +# Description: Functions for character chat cards. +# +# Imports +import json +import logging +import io +import base64 +import time +from typing import Dict, Any, Optional, List, Tuple +# +# External Imports +from PIL import Image +# +# Local imports +from App_Function_Libraries.DB.DB_Manager import get_character_card_by_id, get_character_chat_by_id +from App_Function_Libraries.Metrics.metrics_logger import log_counter, log_histogram +# +# Constants +#################################################################################################### +# +# Functions + +# Using https://github.com/malfoyslastname/character-card-spec-v2 as the standard for v2 character cards + +################################################################################# +# +# Placeholder functions: + +def replace_placeholders(text: str, char_name: str, user_name: str) -> str: + """ + Replace placeholders in the given text with appropriate values. + + Args: + text (str): The text containing placeholders. + char_name (str): The name of the character. + user_name (str): The name of the user. + + Returns: + str: The text with placeholders replaced. + """ + replacements = { + '{{char}}': char_name, + '{{user}}': user_name, + '{{random_user}}': user_name # Assuming random_user is the same as user for simplicity + } + + for placeholder, value in replacements.items(): + text = text.replace(placeholder, value) + + return text + +def replace_user_placeholder(history, user_name): + """ + Replaces all instances of '{{user}}' in the chat history with the actual user name. + + Args: + history (list): The current chat history as a list of tuples (user_message, bot_message). + user_name (str): The name entered by the user. + + Returns: + list: Updated chat history with placeholders replaced. + """ + if not user_name: + user_name = "User" # Default name if none provided + + updated_history = [] + for user_msg, bot_msg in history: + # Replace in user message + if user_msg: + user_msg = user_msg.replace("{{user}}", user_name) + # Replace in bot message + if bot_msg: + bot_msg = bot_msg.replace("{{user}}", user_name) + updated_history.append((user_msg, bot_msg)) + return updated_history + +# +# End of Placeholder functions +################################################################################# + +################################################################################# +# +# Functions for character card processing: + +def extract_character_id(choice: str) -> int: + """Extract the character ID from the dropdown selection string.""" + log_counter("extract_character_id_attempt") + try: + character_id = int(choice.split('(ID: ')[1].rstrip(')')) + log_counter("extract_character_id_success") + return character_id + except Exception as e: + log_counter("extract_character_id_error", labels={"error": str(e)}) + raise + +def load_character_wrapper(character_id: int, user_name: str) -> Tuple[Dict[str, Any], List[Tuple[Optional[str], str]], Optional[Image.Image]]: + """Wrapper function to load character and image using the extracted ID.""" + log_counter("load_character_wrapper_attempt") + start_time = time.time() + try: + char_data, chat_history, img = load_character_and_image(character_id, user_name) + load_duration = time.time() - start_time + log_histogram("load_character_wrapper_duration", load_duration) + log_counter("load_character_wrapper_success") + return char_data, chat_history, img + except Exception as e: + log_counter("load_character_wrapper_error", labels={"error": str(e)}) + raise + +def parse_character_book(book_data: Dict[str, Any]) -> Dict[str, Any]: + """ + Parse the character book data from a V2 character card. + + Args: + book_data (Dict[str, Any]): The raw character book data from the character card. + + Returns: + Dict[str, Any]: The parsed and structured character book data. + """ + parsed_book = { + 'name': book_data.get('name', ''), + 'description': book_data.get('description', ''), + 'scan_depth': book_data.get('scan_depth'), + 'token_budget': book_data.get('token_budget'), + 'recursive_scanning': book_data.get('recursive_scanning', False), + 'extensions': book_data.get('extensions', {}), + 'entries': [] + } + + for entry in book_data.get('entries', []): + parsed_entry = { + 'keys': entry['keys'], + 'content': entry['content'], + 'extensions': entry.get('extensions', {}), + 'enabled': entry['enabled'], + 'insertion_order': entry['insertion_order'], + 'case_sensitive': entry.get('case_sensitive', False), + 'name': entry.get('name', ''), + 'priority': entry.get('priority'), + 'id': entry.get('id'), + 'comment': entry.get('comment', ''), + 'selective': entry.get('selective', False), + 'secondary_keys': entry.get('secondary_keys', []), + 'constant': entry.get('constant', False), + 'position': entry.get('position') + } + parsed_book['entries'].append(parsed_entry) + + return parsed_book + +def load_character_and_image(character_id: int, user_name: str) -> Tuple[Optional[Dict[str, Any]], List[Tuple[Optional[str], str]], Optional[Image.Image]]: + """ + Load a character and its associated image based on the character ID. + + Args: + character_id (int): The ID of the character to load. + user_name (str): The name of the user, used for placeholder replacement. + + Returns: + Tuple[Optional[Dict[str, Any]], List[Tuple[Optional[str], str]], Optional[Image.Image]]: + A tuple containing the character data, chat history, and character image (if available). + """ + log_counter("load_character_and_image_attempt") + start_time = time.time() + try: + char_data = get_character_card_by_id(character_id) + if not char_data: + log_counter("load_character_and_image_no_data") + logging.warning(f"No character data found for ID: {character_id}") + return None, [], None + + # Replace placeholders in character data + for field in ['first_mes', 'mes_example', 'scenario', 'description', 'personality']: + if field in char_data: + char_data[field] = replace_placeholders(char_data[field], char_data['name'], user_name) + + # Replace placeholders in first_mes + first_mes = char_data.get('first_mes', "Hello! I'm ready to chat.") + first_mes = replace_placeholders(first_mes, char_data['name'], user_name) + + chat_history = [(None, first_mes)] if first_mes else [] + + img = None + if char_data.get('image'): + try: + image_data = base64.b64decode(char_data['image']) + img = Image.open(io.BytesIO(image_data)).convert("RGBA") + log_counter("load_character_image_success") + except Exception as e: + log_counter("load_character_image_error", labels={"error": str(e)}) + logging.error(f"Error processing image for character '{char_data['name']}': {e}") + + load_duration = time.time() - start_time + log_histogram("load_character_and_image_duration", load_duration) + log_counter("load_character_and_image_success") + return char_data, chat_history, img + + except Exception as e: + log_counter("load_character_and_image_error", labels={"error": str(e)}) + logging.error(f"Error in load_character_and_image: {e}") + return None, [], None + +def load_chat_and_character(chat_id: int, user_name: str) -> Tuple[Optional[Dict[str, Any]], List[Tuple[str, str]], Optional[Image.Image]]: + """ + Load a chat and its associated character, including the character image and process templates. + + Args: + chat_id (int): The ID of the chat to load. + user_name (str): The name of the user. + + Returns: + Tuple[Optional[Dict[str, Any]], List[Tuple[str, str]], Optional[Image.Image]]: + A tuple containing the character data, processed chat history, and character image (if available). + """ + log_counter("load_chat_and_character_attempt") + start_time = time.time() + try: + # Load the chat + chat = get_character_chat_by_id(chat_id) + if not chat: + log_counter("load_chat_and_character_no_chat") + logging.warning(f"No chat found with ID: {chat_id}") + return None, [], None + + # Load the associated character + character_id = chat['character_id'] + char_data = get_character_card_by_id(character_id) + if not char_data: + log_counter("load_chat_and_character_no_character") + logging.warning(f"No character found for chat ID: {chat_id}") + return None, chat['chat_history'], None + + # Process the chat history + processed_history = process_chat_history(chat['chat_history'], char_data['name'], user_name) + + # Load the character image + img = None + if char_data.get('image'): + try: + image_data = base64.b64decode(char_data['image']) + img = Image.open(io.BytesIO(image_data)).convert("RGBA") + log_counter("load_chat_character_image_success") + except Exception as e: + log_counter("load_chat_character_image_error", labels={"error": str(e)}) + logging.error(f"Error processing image for character '{char_data['name']}': {e}") + + # Process character data templates + for field in ['first_mes', 'mes_example', 'scenario', 'description', 'personality']: + if field in char_data: + char_data[field] = replace_placeholders(char_data[field], char_data['name'], user_name) + + load_duration = time.time() - start_time + log_histogram("load_chat_and_character_duration", load_duration) + log_counter("load_chat_and_character_success") + return char_data, processed_history, img + + except Exception as e: + log_counter("load_chat_and_character_error", labels={"error": str(e)}) + logging.error(f"Error in load_chat_and_character: {e}") + return None, [], None + + +def extract_json_from_image(image_file): + logging.debug(f"Attempting to extract JSON from image: {image_file.name}") + log_counter("extract_json_from_image_attempt") + start_time = time.time() + try: + with Image.open(image_file) as img: + logging.debug("Image opened successfully") + metadata = img.info + if 'chara' in metadata: + logging.debug("Found 'chara' in image metadata") + chara_content = metadata['chara'] + logging.debug(f"Content of 'chara' metadata (first 100 chars): {chara_content[:100]}...") + try: + decoded_content = base64.b64decode(chara_content).decode('utf-8') + logging.debug(f"Decoded content (first 100 chars): {decoded_content[:100]}...") + log_counter("extract_json_from_image_metadata_success") + return decoded_content + except Exception as e: + logging.error(f"Error decoding base64 content: {e}") + log_counter("extract_json_from_image_decode_error", labels={"error": str(e)}) + + logging.warning("'chara' not found in metadata, attempting to find JSON data in image bytes") + # Alternative method to extract embedded JSON from image bytes if metadata is not available + img_byte_arr = io.BytesIO() + img.save(img_byte_arr, format='PNG') + img_bytes = img_byte_arr.getvalue() + img_str = img_bytes.decode('latin1') + + # Search for JSON-like structures in the image bytes + json_start = img_str.find('{') + json_end = img_str.rfind('}') + if json_start != -1 and json_end != -1 and json_end > json_start: + possible_json = img_str[json_start:json_end+1] + try: + json.loads(possible_json) + logging.debug("Found JSON data in image bytes") + log_counter("extract_json_from_image_bytes_success") + return possible_json + except json.JSONDecodeError: + logging.debug("No valid JSON found in image bytes") + log_counter("extract_json_from_image_invalid_json") + + logging.warning("No JSON data found in the image") + log_counter("extract_json_from_image_no_json_found") + except Exception as e: + log_counter("extract_json_from_image_error", labels={"error": str(e)}) + logging.error(f"Error extracting JSON from image: {e}") + + extract_duration = time.time() - start_time + log_histogram("extract_json_from_image_duration", extract_duration) + return None + + +def load_chat_history(file): + log_counter("load_chat_history_attempt") + start_time = time.time() + try: + content = file.read().decode('utf-8') + chat_data = json.loads(content) + + # Extract history and character name from the loaded data + history = chat_data.get('history') or chat_data.get('messages') + character_name = chat_data.get('character') or chat_data.get('character_name') + + if not history or not character_name: + log_counter("load_chat_history_incomplete_data") + logging.error("Chat history or character name missing in the imported file.") + return None, None + + load_duration = time.time() - start_time + log_histogram("load_chat_history_duration", load_duration) + log_counter("load_chat_history_success") + return history, character_name + except Exception as e: + log_counter("load_chat_history_error", labels={"error": str(e)}) + logging.error(f"Error loading chat history: {e}") + return None, None + + +def process_chat_history(chat_history: List[Tuple[str, str]], char_name: str, user_name: str) -> List[Tuple[str, str]]: + """ + Process the chat history to replace placeholders in both user and character messages. + + Args: + chat_history (List[Tuple[str, str]]): The chat history. + char_name (str): The name of the character. + user_name (str): The name of the user. + + Returns: + List[Tuple[str, str]]: The processed chat history. + """ + log_counter("process_chat_history_attempt") + start_time = time.time() + try: + processed_history = [] + for user_msg, char_msg in chat_history: + if user_msg: + user_msg = replace_placeholders(user_msg, char_name, user_name) + if char_msg: + char_msg = replace_placeholders(char_msg, char_name, user_name) + processed_history.append((user_msg, char_msg)) + + process_duration = time.time() - start_time + log_histogram("process_chat_history_duration", process_duration) + log_counter("process_chat_history_success", labels={"message_count": len(chat_history)}) + return processed_history + except Exception as e: + log_counter("process_chat_history_error", labels={"error": str(e)}) + logging.error(f"Error processing chat history: {e}") + raise + +def validate_character_book(book_data): + """ + Validate the 'character_book' field in the character card. + + Args: + book_data (dict): The character book data. + + Returns: + Tuple[bool, List[str]]: A tuple containing a boolean indicating validity and a list of validation messages. + """ + validation_messages = [] + + # Optional fields with expected types + optional_fields = { + 'name': str, + 'description': str, + 'scan_depth': (int, float), + 'token_budget': (int, float), + 'recursive_scanning': bool, + 'extensions': dict, + 'entries': list + } + + for field, expected_type in optional_fields.items(): + if field in book_data: + if not isinstance(book_data[field], expected_type): + validation_messages.append(f"Field 'character_book.{field}' must be of type '{expected_type}'.") + # 'entries' is required + if 'entries' not in book_data or not isinstance(book_data['entries'], list): + validation_messages.append("Field 'character_book.entries' is required and must be a list.") + return False, validation_messages + + # Validate each entry in 'entries' + entries = book_data.get('entries', []) + entry_ids = set() + for idx, entry in enumerate(entries): + is_valid_entry, entry_messages = validate_character_book_entry(entry, idx, entry_ids) + if not is_valid_entry: + validation_messages.extend(entry_messages) + + is_valid = len(validation_messages) == 0 + return is_valid, validation_messages + +def validate_character_book_entry(entry, idx, entry_ids): + """ + Validate an entry in the 'character_book.entries' list. + + Args: + entry (dict): The entry data. + idx (int): The index of the entry in the list. + entry_ids (set): A set of existing entry IDs for uniqueness checking. + + Returns: + Tuple[bool, List[str]]: A tuple containing a boolean indicating validity and a list of validation messages. + """ + validation_messages = [] + required_fields = { + 'keys': list, + 'content': str, + 'extensions': dict, + 'enabled': bool, + 'insertion_order': (int, float) + } + + for field, expected_type in required_fields.items(): + if field not in entry: + validation_messages.append(f"Entry {idx}: Missing required field '{field}'.") + elif not isinstance(entry[field], expected_type): + validation_messages.append(f"Entry {idx}: Field '{field}' must be of type '{expected_type}'.") + elif field == 'content' and not entry[field].strip(): + validation_messages.append(f"Entry {idx}: Field 'content' cannot be empty.") + elif field == 'keys' and not entry[field]: + validation_messages.append(f"Entry {idx}: Field 'keys' cannot be empty.") + + # Optional fields + optional_fields = { + 'case_sensitive': bool, + 'name': str, + 'priority': (int, float), + 'id': (int, float), + 'comment': str, + 'selective': bool, + 'secondary_keys': list, + 'constant': bool, + 'position': str # Should be 'before_char' or 'after_char' + } + + for field, expected_type in optional_fields.items(): + if field in entry and not isinstance(entry[field], expected_type): + validation_messages.append(f"Entry {idx}: Field '{field}' must be of type '{expected_type}'.") + + # Validate 'position' value if present + if 'position' in entry: + if entry['position'] not in ['before_char', 'after_char']: + validation_messages.append(f"Entry {idx}: Field 'position' must be 'before_char' or 'after_char'.") + + # Validate 'secondary_keys' if 'selective' is True + if entry.get('selective', False): + if 'secondary_keys' not in entry or not isinstance(entry['secondary_keys'], list): + validation_messages.append(f"Entry {idx}: 'secondary_keys' must be a list when 'selective' is True.") + elif not entry['secondary_keys']: + validation_messages.append(f"Entry {idx}: 'secondary_keys' cannot be empty when 'selective' is True.") + + # Validate 'keys' list elements + if 'keys' in entry and isinstance(entry['keys'], list): + for i, key in enumerate(entry['keys']): + if not isinstance(key, str) or not key.strip(): + validation_messages.append(f"Entry {idx}: Element {i} in 'keys' must be a non-empty string.") + + # Validate 'secondary_keys' list elements + if 'secondary_keys' in entry and isinstance(entry['secondary_keys'], list): + for i, key in enumerate(entry['secondary_keys']): + if not isinstance(key, str) or not key.strip(): + validation_messages.append(f"Entry {idx}: Element {i} in 'secondary_keys' must be a non-empty string.") + + # Validate 'id' uniqueness + if 'id' in entry: + entry_id = entry['id'] + if entry_id in entry_ids: + validation_messages.append \ + (f"Entry {idx}: Duplicate 'id' value '{entry_id}'. Each entry 'id' must be unique.") + else: + entry_ids.add(entry_id) + + # Validate 'extensions' keys are namespaced + if 'extensions' in entry and isinstance(entry['extensions'], dict): + for key in entry['extensions'].keys(): + if '/' not in key and '_' not in key: + validation_messages.append \ + (f"Entry {idx}: Extension key '{key}' in 'extensions' should be namespaced to prevent conflicts.") + + is_valid = len(validation_messages) == 0 + return is_valid, validation_messages + +def validate_v2_card(card_data): + """ + Validate a character card according to the V2 specification. + + Args: + card_data (dict): The parsed character card data. + + Returns: + Tuple[bool, List[str]]: A tuple containing a boolean indicating validity and a list of validation messages. + """ + validation_messages = [] + + # Check top-level fields + if 'spec' not in card_data: + validation_messages.append("Missing 'spec' field.") + elif card_data['spec'] != 'chara_card_v2': + validation_messages.append(f"Invalid 'spec' value: {card_data['spec']}. Expected 'chara_card_v2'.") + + if 'spec_version' not in card_data: + validation_messages.append("Missing 'spec_version' field.") + else: + # Ensure 'spec_version' is '2.0' or higher + try: + spec_version = float(card_data['spec_version']) + if spec_version < 2.0: + validation_messages.append \ + (f"'spec_version' must be '2.0' or higher. Found '{card_data['spec_version']}'.") + except ValueError: + validation_messages.append \ + (f"Invalid 'spec_version' format: {card_data['spec_version']}. Must be a number as a string.") + + if 'data' not in card_data: + validation_messages.append("Missing 'data' field.") + return False, validation_messages # Cannot proceed without 'data' field + + data = card_data['data'] + + # Required fields in 'data' + required_fields = ['name', 'description', 'personality', 'scenario', 'first_mes', 'mes_example'] + for field in required_fields: + if field not in data: + validation_messages.append(f"Missing required field in 'data': '{field}'.") + elif not isinstance(data[field], str): + validation_messages.append(f"Field '{field}' must be a string.") + elif not data[field].strip(): + validation_messages.append(f"Field '{field}' cannot be empty.") + + # Optional fields with expected types + optional_fields = { + 'creator_notes': str, + 'system_prompt': str, + 'post_history_instructions': str, + 'alternate_greetings': list, + 'tags': list, + 'creator': str, + 'character_version': str, + 'extensions': dict, + 'character_book': dict # If present, should be a dict + } + + for field, expected_type in optional_fields.items(): + if field in data: + if not isinstance(data[field], expected_type): + validation_messages.append(f"Field '{field}' must be of type '{expected_type.__name__}'.") + elif field == 'extensions': + # Validate that extensions keys are properly namespaced + for key in data[field].keys(): + if '/' not in key and '_' not in key: + validation_messages.append \ + (f"Extension key '{key}' in 'extensions' should be namespaced to prevent conflicts.") + + # If 'alternate_greetings' is present, check that it's a list of non-empty strings + if 'alternate_greetings' in data and isinstance(data['alternate_greetings'], list): + for idx, greeting in enumerate(data['alternate_greetings']): + if not isinstance(greeting, str) or not greeting.strip(): + validation_messages.append(f"Element {idx} in 'alternate_greetings' must be a non-empty string.") + + # If 'tags' is present, check that it's a list of non-empty strings + if 'tags' in data and isinstance(data['tags'], list): + for idx, tag in enumerate(data['tags']): + if not isinstance(tag, str) or not tag.strip(): + validation_messages.append(f"Element {idx} in 'tags' must be a non-empty string.") + + # Validate 'extensions' field + if 'extensions' in data and not isinstance(data['extensions'], dict): + validation_messages.append("Field 'extensions' must be a dictionary.") + + # Validate 'character_book' if present + if 'character_book' in data: + is_valid_book, book_messages = validate_character_book(data['character_book']) + if not is_valid_book: + validation_messages.extend(book_messages) + + is_valid = len(validation_messages) == 0 + return is_valid, validation_messages + +# +# End of File +#################################################################################################### diff --git a/App_Function_Libraries/Character_Chat/__init__.py b/App_Function_Libraries/Character_Chat/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/App_Function_Libraries/Chat.py b/App_Function_Libraries/Chat.py new file mode 100644 index 0000000000000000000000000000000000000000..edda250bc09e8f6283b7ae4572e28dfcc2957c7f --- /dev/null +++ b/App_Function_Libraries/Chat.py @@ -0,0 +1,439 @@ +# Chat.py +# Chat functions for interacting with the LLMs as chatbots +import base64 +# Imports +import json +import logging +import os +import re +import tempfile +import time +from datetime import datetime +from pathlib import Path +# +# External Imports +# +# Local Imports +from App_Function_Libraries.DB.DB_Manager import get_conversation_name, save_chat_history_to_database +from App_Function_Libraries.LLM_API_Calls import chat_with_openai, chat_with_anthropic, chat_with_cohere, \ + chat_with_groq, chat_with_openrouter, chat_with_deepseek, chat_with_mistral, chat_with_huggingface +from App_Function_Libraries.LLM_API_Calls_Local import chat_with_aphrodite, chat_with_local_llm, chat_with_ollama, \ + chat_with_kobold, chat_with_llama, chat_with_oobabooga, chat_with_tabbyapi, chat_with_vllm, chat_with_custom_openai +from App_Function_Libraries.DB.SQLite_DB import load_media_content +from App_Function_Libraries.Utils.Utils import generate_unique_filename, load_and_log_configs +from App_Function_Libraries.Metrics.metrics_logger import log_counter, log_histogram +# +#################################################################################################### +# +# Functions: + +def chat_api_call(api_endpoint, api_key, input_data, prompt, temp, system_message=None): + log_counter("chat_api_call_attempt", labels={"api_endpoint": api_endpoint}) + start_time = time.time() + if not api_key: + api_key = None + model = None + try: + logging.info(f"Debug - Chat API Call - API Endpoint: {api_endpoint}") + logging.info(f"Debug - Chat API Call - API Key: {api_key}") + logging.info(f"Debug - Chat chat_api_call - API Endpoint: {api_endpoint}") + if api_endpoint.lower() == 'openai': + response = chat_with_openai(api_key, input_data, prompt, temp, system_message) + + elif api_endpoint.lower() == 'anthropic': + # Retrieve the model from config + loaded_config_data = load_and_log_configs() + model = loaded_config_data['models']['anthropic'] if loaded_config_data else None + response = chat_with_anthropic( + api_key=api_key, + input_data=input_data, + model=model, + custom_prompt_arg=prompt, + system_prompt=system_message + ) + + elif api_endpoint.lower() == "cohere": + response = chat_with_cohere( + api_key, + input_data, + model=model, + custom_prompt_arg=prompt, + system_prompt=system_message, + temp=temp + ) + + elif api_endpoint.lower() == "groq": + response = chat_with_groq(api_key, input_data, prompt, temp, system_message) + + elif api_endpoint.lower() == "openrouter": + response = chat_with_openrouter(api_key, input_data, prompt, temp, system_message) + + elif api_endpoint.lower() == "deepseek": + response = chat_with_deepseek(api_key, input_data, prompt, temp, system_message) + + elif api_endpoint.lower() == "mistral": + response = chat_with_mistral(api_key, input_data, prompt, temp, system_message) + + elif api_endpoint.lower() == "llama.cpp": + response = chat_with_llama(input_data, prompt, temp, None, api_key, system_message) + elif api_endpoint.lower() == "kobold": + response = chat_with_kobold(input_data, api_key, prompt, temp, system_message) + + elif api_endpoint.lower() == "ooba": + response = chat_with_oobabooga(input_data, api_key, prompt, temp, system_message) + + elif api_endpoint.lower() == "tabbyapi": + response = chat_with_tabbyapi(input_data, prompt, temp, system_message) + + elif api_endpoint.lower() == "vllm": + response = chat_with_vllm(input_data, prompt, system_message) + + elif api_endpoint.lower() == "local-llm": + response = chat_with_local_llm(input_data, prompt, temp, system_message) + + elif api_endpoint.lower() == "huggingface": + response = chat_with_huggingface(api_key, input_data, prompt, temp) # , system_message) + + elif api_endpoint.lower() == "ollama": + response = chat_with_ollama(input_data, prompt, None, api_key, temp, system_message) + + elif api_endpoint.lower() == "aphrodite": + response = chat_with_aphrodite(input_data, prompt, temp, system_message) + + elif api_endpoint.lower() == "custom-openai-api": + response = chat_with_custom_openai(api_key, input_data, prompt, temp, system_message) + + else: + raise ValueError(f"Unsupported API endpoint: {api_endpoint}") + + call_duration = time.time() - start_time + log_histogram("chat_api_call_duration", call_duration, labels={"api_endpoint": api_endpoint}) + log_counter("chat_api_call_success", labels={"api_endpoint": api_endpoint}) + return response + + except Exception as e: + log_counter("chat_api_call_error", labels={"api_endpoint": api_endpoint, "error": str(e)}) + logging.error(f"Error in chat function: {str(e)}") + return f"An error occurred: {str(e)}" + + +def chat(message, history, media_content, selected_parts, api_endpoint, api_key, prompt, temperature, + system_message=None): + log_counter("chat_attempt", labels={"api_endpoint": api_endpoint}) + start_time = time.time() + try: + logging.info(f"Debug - Chat Function - Message: {message}") + logging.info(f"Debug - Chat Function - Media Content: {media_content}") + logging.info(f"Debug - Chat Function - Selected Parts: {selected_parts}") + logging.info(f"Debug - Chat Function - API Endpoint: {api_endpoint}") + # logging.info(f"Debug - Chat Function - Prompt: {prompt}") + + # Ensure selected_parts is a list + if not isinstance(selected_parts, (list, tuple)): + selected_parts = [selected_parts] if selected_parts else [] + + # logging.debug(f"Debug - Chat Function - Selected Parts (after check): {selected_parts}") + + # Combine the selected parts of the media content + combined_content = "\n\n".join( + [f"{part.capitalize()}: {media_content.get(part, '')}" for part in selected_parts if part in media_content]) + # Print first 500 chars + # logging.debug(f"Debug - Chat Function - Combined Content: {combined_content[:500]}...") + + # Prepare the input for the API + input_data = f"{combined_content}\n\n" if combined_content else "" + for old_message, old_response in history: + input_data += f"{old_message}\nAssistant: {old_response}\n\n" + input_data += f"{message}\n" + + if system_message: + print(f"System message: {system_message}") + logging.debug(f"Debug - Chat Function - System Message: {system_message}") + temperature = float(temperature) if temperature else 0.7 + temp = temperature + + logging.debug(f"Debug - Chat Function - Temperature: {temperature}") + logging.debug(f"Debug - Chat Function - API Key: {api_key[:10]}") + logging.debug(f"Debug - Chat Function - Prompt: {prompt}") + + # Use the existing API request code based on the selected endpoint + response = chat_api_call(api_endpoint, api_key, input_data, prompt, temp, system_message) + + chat_duration = time.time() - start_time + log_histogram("chat_duration", chat_duration, labels={"api_endpoint": api_endpoint}) + log_counter("chat_success", labels={"api_endpoint": api_endpoint}) + return response + except Exception as e: + log_counter("chat_error", labels={"api_endpoint": api_endpoint, "error": str(e)}) + logging.error(f"Error in chat function: {str(e)}") + return f"An error occurred: {str(e)}" + + +def save_chat_history_to_db_wrapper(chatbot, conversation_id, media_content, media_name=None): + log_counter("save_chat_history_to_db_attempt") + start_time = time.time() + logging.info(f"Attempting to save chat history. Media content type: {type(media_content)}") + try: + # Extract the media_id and media_name from the media_content + media_id = None + if isinstance(media_content, dict): + media_id = None + logging.debug(f"Media content keys: {media_content.keys()}") + if 'content' in media_content: + try: + content = media_content['content'] + if isinstance(content, str): + content_json = json.loads(content) + elif isinstance(content, dict): + content_json = content + else: + raise ValueError(f"Unexpected content type: {type(content)}") + + # Use the webpage_url as the media_id + media_id = content_json.get('webpage_url') + # Use the title as the media_name + media_name = content_json.get('title') + + logging.info(f"Extracted media_id: {media_id}, media_name: {media_name}") + except json.JSONDecodeError: + logging.error("Failed to decode JSON from media_content['content']") + except Exception as e: + logging.error(f"Error processing media_content: {str(e)}") + else: + logging.warning("'content' key not found in media_content") + else: + logging.warning(f"media_content is not a dictionary. Type: {type(media_content)}") + + if media_id is None: + # If we couldn't find a media_id, we'll use a placeholder + media_id = "unknown_media" + logging.warning(f"Unable to extract media_id from media_content. Using placeholder: {media_id}") + + if media_name is None: + media_name = "Unnamed Media" + logging.warning(f"Unable to extract media_name from media_content. Using placeholder: {media_name}") + + # Generate a unique conversation name using media_id and current timestamp + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + conversation_name = f"{media_name}_{timestamp}" + + new_conversation_id = save_chat_history_to_database(chatbot, conversation_id, media_id, media_name, + conversation_name) + save_duration = time.time() - start_time + log_histogram("save_chat_history_to_db_duration", save_duration) + log_counter("save_chat_history_to_db_success") + return new_conversation_id, f"Chat history saved successfully as {conversation_name}!" + except Exception as e: + log_counter("save_chat_history_to_db_error", labels={"error": str(e)}) + error_message = f"Failed to save chat history: {str(e)}" + logging.error(error_message, exc_info=True) + return conversation_id, error_message + + +def save_chat_history(history, conversation_id, media_content): + log_counter("save_chat_history_attempt") + start_time = time.time() + try: + content, conversation_name = generate_chat_history_content(history, conversation_id, media_content) + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + safe_conversation_name = re.sub(r'[^a-zA-Z0-9_-]', '_', conversation_name) + base_filename = f"{safe_conversation_name}_{timestamp}.json" + + # Create a temporary file + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json') as temp_file: + temp_file.write(content) + temp_file_path = temp_file.name + + # Generate a unique filename + unique_filename = generate_unique_filename(os.path.dirname(temp_file_path), base_filename) + final_path = os.path.join(os.path.dirname(temp_file_path), unique_filename) + + # Rename the temporary file to the unique filename + os.rename(temp_file_path, final_path) + + save_duration = time.time() - start_time + log_histogram("save_chat_history_duration", save_duration) + log_counter("save_chat_history_success") + return final_path + except Exception as e: + log_counter("save_chat_history_error", labels={"error": str(e)}) + logging.error(f"Error saving chat history: {str(e)}") + return None + + +def generate_chat_history_content(history, conversation_id, media_content): + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + + conversation_name = get_conversation_name(conversation_id) + + if not conversation_name: + media_name = extract_media_name(media_content) + if media_name: + conversation_name = f"{media_name}-chat" + else: + conversation_name = f"chat-{timestamp}" # Fallback name + + chat_data = { + "conversation_id": conversation_id, + "conversation_name": conversation_name, + "timestamp": timestamp, + "history": [ + { + "role": "user" if i % 2 == 0 else "bot", + "content": msg[0] if isinstance(msg, tuple) else msg + } + for i, msg in enumerate(history) + ] + } + + return json.dumps(chat_data, indent=2), conversation_name + + +def extract_media_name(media_content): + if isinstance(media_content, dict): + content = media_content.get('content', {}) + if isinstance(content, str): + try: + content = json.loads(content) + except json.JSONDecodeError: + logging.warning("Failed to parse media_content JSON string") + return None + + # Try to extract title from the content + if isinstance(content, dict): + return content.get('title') or content.get('name') + + logging.warning(f"Unexpected media_content format: {type(media_content)}") + return None + + +def update_chat_content(selected_item, use_content, use_summary, use_prompt, item_mapping): + log_counter("update_chat_content_attempt") + start_time = time.time() + logging.debug(f"Debug - Update Chat Content - Selected Item: {selected_item}\n") + logging.debug(f"Debug - Update Chat Content - Use Content: {use_content}\n\n\n\n") + logging.debug(f"Debug - Update Chat Content - Use Summary: {use_summary}\n\n") + logging.debug(f"Debug - Update Chat Content - Use Prompt: {use_prompt}\n\n") + logging.debug(f"Debug - Update Chat Content - Item Mapping: {item_mapping}\n\n") + + if selected_item and selected_item in item_mapping: + media_id = item_mapping[selected_item] + content = load_media_content(media_id) + selected_parts = [] + if use_content and "content" in content: + selected_parts.append("content") + if use_summary and "summary" in content: + selected_parts.append("summary") + if use_prompt and "prompt" in content: + selected_parts.append("prompt") + + # Modified debug print + if isinstance(content, dict): + print(f"Debug - Update Chat Content - Content keys: {list(content.keys())}") + for key, value in content.items(): + print(f"Debug - Update Chat Content - {key} (first 500 char): {str(value)[:500]}\n\n\n\n") + else: + print(f"Debug - Update Chat Content - Content(first 500 char): {str(content)[:500]}\n\n\n\n") + + print(f"Debug - Update Chat Content - Selected Parts: {selected_parts}") + update_duration = time.time() - start_time + log_histogram("update_chat_content_duration", update_duration) + log_counter("update_chat_content_success") + return content, selected_parts + else: + log_counter("update_chat_content_error", labels={"error": str("No item selected or item not in mapping")}) + print(f"Debug - Update Chat Content - No item selected or item not in mapping") + return {}, [] + +# +# End of Chat functions +####################################################################################################################### + + +####################################################################################################################### +# +# Character Card Functions + +CHARACTERS_FILE = Path('.', 'Helper_Scripts', 'Character_Cards', 'Characters.json') + + +def save_character(character_data): + log_counter("save_character_attempt") + start_time = time.time() + characters_file = os.path.join(os.path.dirname(__file__), '..', 'Helper_Scripts', 'Character_Cards', 'Characters.json') + characters_dir = os.path.dirname(characters_file) + + try: + if os.path.exists(characters_file): + with open(characters_file, 'r') as f: + characters = json.load(f) + else: + characters = {} + + char_name = character_data['name'] + + # Save the image separately if it exists + if 'image' in character_data: + img_data = base64.b64decode(character_data['image']) + img_filename = f"{char_name.replace(' ', '_')}.png" + img_path = os.path.join(characters_dir, img_filename) + with open(img_path, 'wb') as f: + f.write(img_data) + character_data['image_path'] = os.path.abspath(img_path) + del character_data['image'] # Remove the base64 image data from the JSON + + characters[char_name] = character_data + + with open(characters_file, 'w') as f: + json.dump(characters, f, indent=2) + + save_duration = time.time() - start_time + log_histogram("save_character_duration", save_duration) + log_counter("save_character_success") + logging.info(f"Character '{char_name}' saved successfully.") + except Exception as e: + log_counter("save_character_error", labels={"error": str(e)}) + logging.error(f"Error saving character: {str(e)}") + + +def load_characters(): + log_counter("load_characters_attempt") + start_time = time.time() + try: + characters_file = os.path.join(os.path.dirname(__file__), '..', 'Helper_Scripts', 'Character_Cards', 'Characters.json') + if os.path.exists(characters_file): + with open(characters_file, 'r') as f: + characters = json.load(f) + logging.debug(f"Loaded {len(characters)} characters from {characters_file}") + load_duration = time.time() - start_time + log_histogram("load_characters_duration", load_duration) + log_counter("load_characters_success", labels={"character_count": len(characters)}) + return characters + else: + logging.warning(f"Characters file not found: {characters_file}") + return {} + except Exception as e: + log_counter("load_characters_error", labels={"error": str(e)}) + return {} + + + +def get_character_names(): + log_counter("get_character_names_attempt") + start_time = time.time() + try: + characters = load_characters() + names = list(characters.keys()) + get_names_duration = time.time() - start_time + log_histogram("get_character_names_duration", get_names_duration) + log_counter("get_character_names_success", labels={"name_count": len(names)}) + return names + except Exception as e: + log_counter("get_character_names_error", labels={"error": str(e)}) + logging.error(f"Error getting character names: {str(e)}") + return [] + +# +# End of Chat.py +########################################################################################################################## diff --git a/App_Function_Libraries/Chunk_Lib.py b/App_Function_Libraries/Chunk_Lib.py new file mode 100644 index 0000000000000000000000000000000000000000..f60bcf2e6f450c46653f428e513a85fd4f4564dd --- /dev/null +++ b/App_Function_Libraries/Chunk_Lib.py @@ -0,0 +1,1051 @@ +# Chunk_Lib.py +######################################### +# Chunking Library +# This library is used to perform chunking of input files. +# Currently, uses naive approaches. Nothing fancy. +# +#### +# Import necessary libraries +import hashlib +import json +import logging +import re +from typing import Any, Dict, List, Optional, Tuple +# +# Import 3rd party +from openai import OpenAI +from tqdm import tqdm +from langdetect import detect +from transformers import GPT2Tokenizer +import nltk +from nltk.tokenize import sent_tokenize, word_tokenize +from sklearn.feature_extraction.text import TfidfVectorizer +from sklearn.metrics.pairwise import cosine_similarity +# +# Import Local +from App_Function_Libraries.Tokenization_Methods_Lib import openai_tokenize +from App_Function_Libraries.Utils.Utils import load_comprehensive_config +# +####################################################################################################################### +# Config Settings +# +# +# FIXME - Make sure it only downloads if it already exists, and does a check first. +# Ensure NLTK data is downloaded +def ensure_nltk_data(): + try: + nltk.data.find('tokenizers/punkt') + except LookupError: + nltk.download('punkt') +ensure_nltk_data() + +# +# Load GPT2 tokenizer +tokenizer = GPT2Tokenizer.from_pretrained("gpt2") +# +# Load configuration +config = load_comprehensive_config() +# Embedding Chunking options +chunk_options = { + 'method': config.get('Chunking', 'method', fallback='words'), + 'max_size': config.getint('Chunking', 'max_size', fallback=400), + 'overlap': config.getint('Chunking', 'overlap', fallback=200), + 'adaptive': config.getboolean('Chunking', 'adaptive', fallback=False), + 'multi_level': config.getboolean('Chunking', 'multi_level', fallback=False), + 'language': config.get('Chunking', 'language', fallback='english') +} + +openai_api_key = config.get('API', 'openai_api_key') +# +# End of settings +####################################################################################################################### +# +# Functions: + +# Create a chunking class for refactoring FIXME +# class Chunker: +# def __init__(self, tokenizer: GPT2Tokenizer): +# self.tokenizer = tokenizer +# +# def detect_language(self, text: str) -> str: +# try: +# return detect(text) +# except: +# return 'en' +# +# def chunk_text(self, text: str, method: str, max_size: int, overlap: int, language: str = None) -> List[str]: +# if language is None: +# language = self.detect_language(text) +# +# if method == 'words': +# return self.chunk_text_by_words(text, max_size, overlap, language) +# elif method == 'sentences': +# return self.chunk_text_by_sentences(text, max_size, overlap, language) +# elif method == 'paragraphs': +# return self.chunk_text_by_paragraphs(text, max_size, overlap) +# elif method == 'tokens': +# return self.chunk_text_by_tokens(text, max_size, overlap, language) +# elif method == 'semantic': +# return self.semantic_chunking(text, max_size) +# else: +# return [text] + +def detect_language(text: str) -> str: + try: + return detect(text) + except: + # Default to English if detection fails + return 'en' + + +def load_document(file_path: str) -> str: + with open(file_path, 'r', encoding='utf-8') as file: + text = file.read() + return re.sub(r'\s+', ' ', text).strip() + + +def improved_chunking_process(text: str, chunk_options: Dict[str, Any] = None) -> List[Dict[str, Any]]: + logging.debug("Improved chunking process started...") + + # Extract JSON metadata if present + json_content = {} + try: + json_end = text.index("}\n") + 1 + json_content = json.loads(text[:json_end]) + text = text[json_end:].strip() + logging.debug(f"Extracted JSON metadata: {json_content}") + except (ValueError, json.JSONDecodeError): + logging.debug("No JSON metadata found at the beginning of the text") + + # Extract any additional header text + header_match = re.match(r"(This text was transcribed using.*?)\n\n", text, re.DOTALL) + header_text = "" + if header_match: + header_text = header_match.group(1) + text = text[len(header_text):].strip() + logging.debug(f"Extracted header text: {header_text}") + + options = chunk_options.copy() if chunk_options else {} + if chunk_options: + options.update(chunk_options) + + chunk_method = options.get('method', 'words') + max_size = options.get('max_size', 2000) + overlap = options.get('overlap', 0) + language = options.get('language', None) + + if language is None: + language = detect_language(text) + + if chunk_method == 'json': + chunks = chunk_text_by_json(text, max_size=max_size, overlap=overlap) + else: + chunks = chunk_text(text, chunk_method, max_size, overlap, language) + + chunks_with_metadata = [] + total_chunks = len(chunks) + for i, chunk in enumerate(chunks): + metadata = { + 'chunk_index': i + 1, + 'total_chunks': total_chunks, + 'chunk_method': chunk_method, + 'max_size': max_size, + 'overlap': overlap, + 'language': language, + 'relative_position': (i + 1) / total_chunks + } + metadata.update(json_content) # Add the extracted JSON content to metadata + metadata['header_text'] = header_text # Add the header text to metadata + + if chunk_method == 'json': + chunk_text_content = json.dumps(chunk['json'], ensure_ascii=False) + else: + chunk_text_content = chunk + + chunks_with_metadata.append({ + 'text': chunk_text_content, + 'metadata': metadata + }) + + return chunks_with_metadata + + +def multi_level_chunking(text: str, method: str, max_size: int, overlap: int, language: str) -> List[str]: + logging.debug("Multi-level chunking process started...") + # First level: chunk by paragraphs + paragraphs = chunk_text_by_paragraphs(text, max_size * 2, overlap) + + # Second level: chunk each paragraph further + chunks = [] + for para in paragraphs: + if method == 'words': + chunks.extend(chunk_text_by_words(para, max_words=max_size, overlap=overlap, language=language)) + elif method == 'sentences': + chunks.extend(chunk_text_by_sentences(para, max_sentences=max_size, overlap=overlap, language=language)) + else: + chunks.append(para) + + return chunks + + +# FIXME - ensure language detection occurs in each chunk function +def chunk_text(text: str, method: str, max_size: int, overlap: int, language: str = None) -> List[str]: + if method == 'words': + logging.debug("Chunking by words...") + return chunk_text_by_words(text, max_words=max_size, overlap=overlap, language=language) + elif method == 'sentences': + logging.debug("Chunking by sentences...") + return chunk_text_by_sentences(text, max_sentences=max_size, overlap=overlap, language=language) + elif method == 'paragraphs': + logging.debug("Chunking by paragraphs...") + return chunk_text_by_paragraphs(text, max_paragraphs=max_size, overlap=overlap) + elif method == 'tokens': + logging.debug("Chunking by tokens...") + return chunk_text_by_tokens(text, max_tokens=max_size, overlap=overlap) + elif method == 'semantic': + logging.debug("Chunking by semantic similarity...") + return semantic_chunking(text, max_chunk_size=max_size) + else: + logging.warning(f"Unknown chunking method '{method}'. Returning full text as a single chunk.") + return [text] + +def determine_chunk_position(relative_position: float) -> str: + if relative_position < 0.33: + return "This chunk is from the beginning of the document" + elif relative_position < 0.66: + return "This chunk is from the middle of the document" + else: + return "This chunk is from the end of the document" + + +def chunk_text_by_words(text: str, max_words: int = 300, overlap: int = 0, language: str = None) -> List[str]: + logging.debug("chunk_text_by_words...") + if language is None: + language = detect_language(text) + + if language.startswith('zh'): # Chinese + import jieba + words = list(jieba.cut(text)) + elif language == 'ja': # Japanese + import fugashi + tagger = fugashi.Tagger() + words = [word.surface for word in tagger(text)] + else: # Default to simple splitting for other languages + words = text.split() + + chunks = [] + for i in range(0, len(words), max_words - overlap): + chunk = ' '.join(words[i:i + max_words]) + chunks.append(chunk) + return post_process_chunks(chunks) + + +def chunk_text_by_sentences(text: str, max_sentences: int = 10, overlap: int = 0, language: str = None) -> List[str]: + logging.debug("chunk_text_by_sentences...") + if language is None: + language = detect_language(text) + + if language.startswith('zh'): # Chinese + import jieba + # Use jieba to perform sentence segmentation + # jieba does not support sentence segmentation out of the box + # Use punctuation as delimiters + sentences = re.split(r'[。！？；]', text) + sentences = [s.strip() for s in sentences if s.strip()] + elif language == 'ja': # Japanese + import fugashi + tagger = fugashi.Tagger() + # Simple sentence segmentation based on punctuation + sentences = re.split(r'[。！？]', text) + sentences = [s.strip() for s in sentences if s.strip()] + else: # Default to NLTK for other languages + try: + sentences = sent_tokenize(text, language=language) + except LookupError: + logging.warning(f"Punkt tokenizer not found for language '{language}'. Using default 'english'.") + sentences = sent_tokenize(text, language='english') + + chunks = [] + previous_overlap = [] + + for i in range(0, len(sentences), max_sentences - overlap): + current_sentences = sentences[i:i + max_sentences] + if overlap > 0 and previous_overlap: + current_sentences = previous_overlap + current_sentences + chunk = ' '.join(current_sentences) + chunks.append(chunk) + previous_overlap = sentences[i + max_sentences - overlap:i + max_sentences] if overlap > 0 else [] + + return post_process_chunks(chunks) + + +def chunk_text_by_paragraphs(text: str, max_paragraphs: int = 5, overlap: int = 0) -> List[str]: + logging.debug("chunk_text_by_paragraphs...") + paragraphs = re.split(r'\n\s*\n', text) + chunks = [] + for i in range(0, len(paragraphs), max_paragraphs - overlap): + chunk = '\n\n'.join(paragraphs[i:i + max_paragraphs]) + chunks.append(chunk) + return post_process_chunks(chunks) + + +def chunk_text_by_tokens(text: str, max_tokens: int = 1000, overlap: int = 0) -> List[str]: + logging.debug("chunk_text_by_tokens...") + # This is a simplified token-based chunking. For more accurate tokenization, + # consider using a proper tokenizer like GPT-2 TokenizerFast + words = text.split() + chunks = [] + current_chunk = [] + current_token_count = 0 + + for word in words: + word_token_count = len(word) // 4 + 1 # Rough estimate of token count + if current_token_count + word_token_count > max_tokens and current_chunk: + chunks.append(' '.join(current_chunk)) + current_chunk = current_chunk[-overlap:] if overlap > 0 else [] + current_token_count = sum(len(w) // 4 + 1 for w in current_chunk) + + current_chunk.append(word) + current_token_count += word_token_count + + if current_chunk: + chunks.append(' '.join(current_chunk)) + + return post_process_chunks(chunks) +# def chunk_text_by_tokens(text: str, max_tokens: int = 1000, overlap: int = 0) -> List[str]: +# logging.debug("chunk_text_by_tokens...") +# # Use GPT2 tokenizer for tokenization +# tokens = tokenizer.encode(text) +# chunks = [] +# for i in range(0, len(tokens), max_tokens - overlap): +# chunk_tokens = tokens[i:i + max_tokens] +# chunk = tokenizer.decode(chunk_tokens) +# chunks.append(chunk) +# return post_process_chunks(chunks) + + +def post_process_chunks(chunks: List[str]) -> List[str]: + return [chunk.strip() for chunk in chunks if chunk.strip()] + + +# FIXME - F +def get_chunk_metadata(chunk: str, full_text: str, chunk_type: str = "generic", + chapter_number: Optional[int] = None, + chapter_pattern: Optional[str] = None, + language: str = None) -> Dict[str, Any]: + """ + Generate metadata for a chunk based on its position in the full text. + """ + chunk_length = len(chunk) + start_index = full_text.find(chunk) + end_index = start_index + chunk_length if start_index != -1 else None + + # Calculate a hash for the chunk + chunk_hash = hashlib.md5(chunk.encode()).hexdigest() + + metadata = { + 'start_index': start_index, + 'end_index': end_index, + 'word_count': len(chunk.split()), + 'char_count': chunk_length, + 'chunk_type': chunk_type, + 'language': language, + 'chunk_hash': chunk_hash, + 'relative_position': start_index / len(full_text) if len(full_text) > 0 and start_index != -1 else 0 + } + + if chunk_type == "chapter": + metadata['chapter_number'] = chapter_number + metadata['chapter_pattern'] = chapter_pattern + + return metadata + + +def process_document_with_metadata(text: str, chunk_options: Dict[str, Any], + document_metadata: Dict[str, Any]) -> Dict[str, Any]: + chunks = improved_chunking_process(text, chunk_options) + + return { + 'document_metadata': document_metadata, + 'chunks': chunks + } + + +# Hybrid approach, chunk each sentence while ensuring total token size does not exceed a maximum number +def chunk_text_hybrid(text: str, max_tokens: int = 1000, overlap: int = 0) -> List[str]: + logging.debug("chunk_text_hybrid...") + sentences = sent_tokenize(text) + chunks = [] + current_chunk = [] + current_length = 0 + + for sentence in sentences: + tokens = tokenizer.encode(sentence) + if current_length + len(tokens) > max_tokens and current_chunk: + chunks.append(' '.join(current_chunk)) + # Handle overlap + if overlap > 0: + overlap_tokens = tokenizer.encode(' '.join(current_chunk[-overlap:])) + current_chunk = current_chunk[-overlap:] + current_length = len(overlap_tokens) + else: + current_chunk = [] + current_length = 0 + + current_chunk.append(sentence) + current_length += len(tokens) + + if current_chunk: + chunks.append(' '.join(current_chunk)) + + return post_process_chunks(chunks) + + +# Thanks openai +def chunk_on_delimiter(input_string: str, + max_tokens: int, + delimiter: str) -> List[str]: + logging.debug("chunk_on_delimiter...") + chunks = input_string.split(delimiter) + combined_chunks, _, dropped_chunk_count = combine_chunks_with_no_minimum( + chunks, max_tokens, chunk_delimiter=delimiter, add_ellipsis_for_overflow=True) + if dropped_chunk_count > 0: + logging.warning(f"Warning: {dropped_chunk_count} chunks were dropped due to exceeding the token limit.") + combined_chunks = [f"{chunk}{delimiter}" for chunk in combined_chunks] + return combined_chunks + + + + +# FIXME +def recursive_summarize_chunks(chunks: List[str], summarize_func, custom_prompt: Optional[str] = None, + temp: Optional[float] = None, system_prompt: Optional[str] = None) -> List[str]: + logging.debug("recursive_summarize_chunks...") + summarized_chunks = [] + current_summary = "" + + logging.debug(f"Summarizing {len(chunks)} chunks recursively...") + logging.debug(f"Temperature is set to {temp}") + for i, chunk in enumerate(chunks): + if i == 0: + current_summary = summarize_func(chunk, custom_prompt, temp, system_prompt) + else: + combined_text = current_summary + "\n\n" + chunk + current_summary = summarize_func(combined_text, custom_prompt, temp, system_prompt) + + summarized_chunks.append(current_summary) + + return summarized_chunks + + +# Sample text for testing +sample_text = """ +Natural language processing (NLP) is a subfield of linguistics, computer science, and artificial intelligence +concerned with the interactions between computers and human language, in particular how to program computers +to process and analyze large amounts of natural language data. The result is a computer capable of "understanding" +the contents of documents, including the contextual nuances of the language within them. The technology can then +accurately extract information and insights contained in the documents as well as categorize and organize the documents themselves. + +Challenges in natural language processing frequently involve speech recognition, natural language understanding, +and natural language generation. + +Natural language processing has its roots in the 1950s. Already in 1950, Alan Turing published an article titled +"Computing Machinery and Intelligence" which proposed what is now called the Turing test as a criterion of intelligence. +""" + +# Example usage of different chunking methods +# print("Chunking by words:") +# print(chunk_text_by_words(sample_text, max_words=50)) +# +# print("\nChunking by sentences:") +# print(chunk_text_by_sentences(sample_text, max_sentences=2)) +# +# print("\nChunking by paragraphs:") +# print(chunk_text_by_paragraphs(sample_text, max_paragraphs=1)) +# +# print("\nChunking by tokens:") +# print(chunk_text_by_tokens(sample_text, max_tokens=50)) +# +# print("\nHybrid chunking:") +# print(chunk_text_hybrid(sample_text, max_tokens=50)) + + + +####################################################################################################################### +# +# Experimental Semantic Chunking +# + +# Chunk text into segments based on semantic similarity +def count_units(text: str, unit: str = 'words') -> int: + if unit == 'words': + return len(text.split()) + elif unit == 'tokens': + return len(tokenizer.encode(text)) + elif unit == 'characters': + return len(text) + else: + raise ValueError("Invalid unit. Choose 'words', 'tokens', or 'characters'.") + + + +def semantic_chunking(text: str, max_chunk_size: int = 2000, unit: str = 'words') -> List[str]: + logging.debug("semantic_chunking...") + sentences = sent_tokenize(text) + vectorizer = TfidfVectorizer() + sentence_vectors = vectorizer.fit_transform(sentences) + + chunks = [] + current_chunk = [] + current_size = 0 + + for i, sentence in enumerate(sentences): + sentence_size = count_units(sentence, unit) + if current_size + sentence_size > max_chunk_size and current_chunk: + chunks.append(' '.join(current_chunk)) + # Use last 3 sentences for overlap + current_chunk = current_chunk[-3:] + current_size = count_units(' '.join(current_chunk), unit) + + current_chunk.append(sentence) + current_size += sentence_size + + if i + 1 < len(sentences): + current_vector = sentence_vectors[i] + next_vector = sentence_vectors[i + 1] + similarity = cosine_similarity(current_vector, next_vector)[0][0] + if similarity < 0.5 and current_size >= max_chunk_size // 2: + chunks.append(' '.join(current_chunk)) + current_chunk = current_chunk[-3:] + current_size = count_units(' '.join(current_chunk), unit) + + if current_chunk: + chunks.append(' '.join(current_chunk)) + + return chunks + + +def semantic_chunk_long_file(file_path: str, max_chunk_size: int = 1000, overlap: int = 100, unit: str = 'words') -> Optional[List[str]]: + logging.debug("semantic_chunk_long_file...") + try: + with open(file_path, 'r', encoding='utf-8') as file: + content = file.read() + + chunks = semantic_chunking(content, max_chunk_size, unit) + return chunks + except Exception as e: + logging.error(f"Error chunking text file: {str(e)}") + return None + +# +# +####################################################################################################################### + + +####################################################################################################################### +# +# Embedding Chunking + +def chunk_for_embedding(text: str, file_name: str, custom_chunk_options: Dict[str, Any] = None) -> List[Dict[str, Any]]: + options = chunk_options.copy() + if custom_chunk_options: + options.update(custom_chunk_options) + + logging.info(f"Chunking options: {options}") + chunks = improved_chunking_process(text, options) + total_chunks = len(chunks) + logging.info(f"Total chunks created: {total_chunks}") + + chunked_text_with_headers = [] + for i, chunk in enumerate(chunks, 1): + chunk_text = chunk['text'] + chunk_position = determine_chunk_position(chunk['metadata']['relative_position']) + chunk_header = f""" + Original Document: {file_name} + Chunk: {i} of {total_chunks} + Position: {chunk_position} + + --- Chunk Content --- + """ + + full_chunk_text = chunk_header + chunk_text + chunk['text'] = full_chunk_text + chunk['metadata']['file_name'] = file_name + chunked_text_with_headers.append(chunk) + + return chunked_text_with_headers + +# +# End of Embedding Chunking +####################################################################################################################### + + +####################################################################################################################### +# +# JSON Chunking + +# FIXME +def chunk_text_by_json(text: str, max_size: int = 1000, overlap: int = 0) -> List[Dict[str, Any]]: + """ + Chunk JSON-formatted text into smaller JSON chunks while preserving structure. + + Parameters: + - text (str): The JSON-formatted text to be chunked. + - max_size (int): Maximum number of items or characters per chunk. + - overlap (int): Number of items or characters to overlap between chunks. + + Returns: + - List[Dict[str, Any]]: A list of chunks with their metadata. + """ + logging.debug("chunk_text_by_json started...") + try: + json_data = json.loads(text) + except json.JSONDecodeError as e: + logging.error(f"Invalid JSON data: {e}") + raise ValueError(f"Invalid JSON data: {e}") + + # Determine if JSON data is a list or a dict + if isinstance(json_data, list): + return chunk_json_list(json_data, max_size, overlap) + elif isinstance(json_data, dict): + return chunk_json_dict(json_data, max_size, overlap) + else: + logging.error("Unsupported JSON structure. Only JSON objects and arrays are supported.") + raise ValueError("Unsupported JSON structure. Only JSON objects and arrays are supported.") + + +def chunk_json_list(json_list: List[Any], max_size: int, overlap: int) -> List[Dict[str, Any]]: + """ + Chunk a JSON array into smaller chunks. + + Parameters: + - json_list (List[Any]): The JSON array to be chunked. + - max_size (int): Maximum number of items per chunk. + - overlap (int): Number of items to overlap between chunks. + + Returns: + - List[Dict[str, Any]]: A list of JSON chunks with metadata. + """ + logging.debug("chunk_json_list started...") + chunks = [] + total_items = len(json_list) + step = max_size - overlap + if step <= 0: + raise ValueError("max_size must be greater than overlap.") + + for i in range(0, total_items, step): + chunk = json_list[i:i + max_size] + metadata = { + 'chunk_index': i // step + 1, + 'total_chunks': (total_items + step - 1) // step, + 'chunk_method': 'json_list', + 'max_size': max_size, + 'overlap': overlap, + 'relative_position': i / total_items + } + chunks.append({ + 'json': chunk, + 'metadata': metadata + }) + + logging.debug(f"chunk_json_list created {len(chunks)} chunks.") + return chunks + + + +def chunk_json_dict(json_dict: Dict[str, Any], max_size: int, overlap: int) -> List[Dict[str, Any]]: + """ + Chunk a JSON object into smaller chunks based on its 'data' key while preserving other keys like 'metadata'. + + Parameters: + - json_dict (Dict[str, Any]): The JSON object to be chunked. + - max_size (int): Maximum number of key-value pairs per chunk in the 'data' section. + - overlap (int): Number of key-value pairs to overlap between chunks. + + Returns: + - List[Dict[str, Any]]: A list of JSON chunks with metadata. + """ + logging.debug("chunk_json_dict started...") + + # Preserve non-chunked sections + preserved_keys = ['metadata'] + preserved_data = {key: value for key, value in json_dict.items() if key in preserved_keys} + + # Identify the chunkable section + chunkable_key = 'data' + if chunkable_key not in json_dict or not isinstance(json_dict[chunkable_key], dict): + logging.error("No chunkable 'data' section found in JSON dictionary.") + raise ValueError("No chunkable 'data' section found in JSON dictionary.") + + chunkable_data = json_dict[chunkable_key] + data_keys = list(chunkable_data.keys()) + total_keys = len(data_keys) + chunks = [] + step = max_size - overlap + if step <= 0: + raise ValueError("max_size must be greater than overlap.") + + # Adjust the loop to prevent creating an extra chunk + for i in range(0, total_keys, step): + chunk_keys = data_keys[i:i + max_size] + + # Handle overlap + if i != 0 and overlap > 0: + overlap_keys = data_keys[i - overlap:i] + chunk_keys = overlap_keys + chunk_keys + + # Remove duplicate keys caused by overlap + unique_chunk_keys = [] + seen_keys = set() + for key in chunk_keys: + if key not in seen_keys: + unique_chunk_keys.append(key) + seen_keys.add(key) + + chunk_data = {key: chunkable_data[key] for key in unique_chunk_keys} + + metadata = { + 'chunk_index': (i // step) + 1, + 'total_chunks': (total_keys + step - 1) // step, + 'chunk_method': 'json_dict', + 'max_size': max_size, + 'overlap': overlap, + 'language': 'english', # Assuming English; modify as needed + 'relative_position': (i // step + 1) / ((total_keys + step - 1) // step) + } + + # Merge preserved data into metadata + metadata.update(preserved_data.get('metadata', {})) + + # Create the chunk with preserved data + chunk = { + 'metadata': preserved_data, + 'data': chunk_data + } + + chunks.append({ + 'json': chunk, + 'metadata': metadata + }) + + logging.debug(f"chunk_json_dict created {len(chunks)} chunks.") + return chunks + + +# +# End of JSON Chunking +####################################################################################################################### + +####################################################################################################################### +# +# OpenAI Rolling Summarization +# + +client = OpenAI(api_key=openai_api_key) +def get_chat_completion(messages, model='gpt-4-turbo'): + response = client.chat.completions.create( + model=model, + messages=messages, + temperature=0, + ) + return response.choices[0].message.content + + +# This function combines text chunks into larger blocks without exceeding a specified token count. +# It returns the combined chunks, their original indices, and the number of dropped chunks due to overflow. +def combine_chunks_with_no_minimum( + chunks: List[str], + max_tokens: int, + chunk_delimiter: str = "\n\n", + header: Optional[str] = None, + add_ellipsis_for_overflow: bool = False, +) -> Tuple[List[str], List[List[int]], int]: + dropped_chunk_count = 0 + output = [] # list to hold the final combined chunks + output_indices = [] # list to hold the indices of the final combined chunks + candidate = [header] if header else [] # list to hold the current combined chunk candidate + candidate_indices = [] + for chunk_i, chunk in enumerate(chunks): + chunk_with_header = [chunk] if not header else [header, chunk] + combined_text = chunk_delimiter.join(candidate + chunk_with_header) + token_count = len(tokenizer.encode(combined_text)) + if token_count > max_tokens: + if add_ellipsis_for_overflow and len(candidate) > 0: + ellipsis_text = chunk_delimiter.join(candidate + ["..."]) + if len(tokenizer.encode(ellipsis_text)) <= max_tokens: + candidate = candidate + ["..."] + dropped_chunk_count += 1 + if len(candidate) > 0: + output.append(chunk_delimiter.join(candidate)) + output_indices.append(candidate_indices) + candidate = chunk_with_header + candidate_indices = [chunk_i] + else: + logging.warning(f"Single chunk at index {chunk_i} exceeds max_tokens and will be dropped.") + dropped_chunk_count += 1 + else: + candidate.extend(chunk_with_header) + candidate_indices.append(chunk_i) + + if candidate: + output.append(chunk_delimiter.join(candidate)) + output_indices.append(candidate_indices) + return output, output_indices, dropped_chunk_count + + +def rolling_summarize(text: str, + detail: float = 0, + model: str = 'gpt-4o', + additional_instructions: Optional[str] = None, + minimum_chunk_size: Optional[int] = 500, + chunk_delimiter: str = ".", + summarize_recursively: bool = False, + verbose: bool = False) -> str: + """ + Summarizes a given text by splitting it into chunks, each of which is summarized individually. + The level of detail in the summary can be adjusted, and the process can optionally be made recursive. + + Parameters: + - text (str): The text to be summarized. + - detail (float, optional): A value between 0 and 1 indicating the desired level of detail in the summary. + - additional_instructions (Optional[str], optional): Additional instructions for the model. + - minimum_chunk_size (Optional[int], optional): The minimum size for text chunks. + - chunk_delimiter (str, optional): The delimiter used to split the text into chunks. + - summarize_recursively (bool, optional): If True, summaries are generated recursively. + - verbose (bool, optional): If True, prints detailed information about the chunking process. + + Returns: + - str: The final compiled summary of the text. + + The function first determines the number of chunks by interpolating between a minimum and a maximum chunk count + based on the `detail` parameter. It then splits the text into chunks and summarizes each chunk. If + `summarize_recursively` is True, each summary is based on the previous summaries, adding more context to the + summarization process. The function returns a compiled summary of all chunks. + """ + + # Check detail is set correctly + assert 0 <= detail <= 1, "Detail must be between 0 and 1." + + # Interpolate the number of chunks based on the detail parameter + text_length = len(tokenizer.encode(text)) + max_chunks = text_length // minimum_chunk_size if minimum_chunk_size else 10 + min_chunks = 1 + num_chunks = int(min_chunks + detail * (max_chunks - min_chunks)) + + # Adjust chunk_size based on interpolated number of chunks + chunk_size = max(minimum_chunk_size, text_length // num_chunks) if num_chunks else text_length + text_chunks = chunk_on_delimiter(text, chunk_size, chunk_delimiter) + if verbose: + print(f"Splitting the text into {len(text_chunks)} chunks to be summarized.") + print(f"Chunk lengths are {[len(tokenizer.encode(x)) for x in text_chunks]} tokens.") + + # Set system message + system_message_content = "Rewrite this text in summarized form." + if additional_instructions: + system_message_content += f"\n\n{additional_instructions}" + + accumulated_summaries = [] + for i, chunk in enumerate(tqdm(text_chunks, desc="Summarizing chunks")): + if summarize_recursively and accumulated_summaries: + # Combine previous summary with current chunk for recursive summarization + combined_text = accumulated_summaries[-1] + "\n\n" + chunk + user_message_content = f"Previous summary and new content to summarize:\n\n{combined_text}" + else: + user_message_content = chunk + + messages = [ + {"role": "system", "content": system_message_content}, + {"role": "user", "content": user_message_content} + ] + + response = get_chat_completion(messages, model=model) + accumulated_summaries.append(response) + + final_summary = '\n\n'.join(accumulated_summaries) + return final_summary + +# +# +####################################################################################################################### +# +# Ebook Chapter Chunking + + +def chunk_ebook_by_chapters(text: str, chunk_options: Dict[str, Any]) -> List[Dict[str, Any]]: + logging.debug("chunk_ebook_by_chapters") + max_chunk_size = int(chunk_options.get('max_size', 300)) + overlap = int(chunk_options.get('overlap', 0)) + custom_pattern = chunk_options.get('custom_chapter_pattern', None) + + # List of chapter heading patterns to try, in order + chapter_patterns = [ + custom_pattern, + r'^#{1,2}\s+', # Markdown style: '# ' or '## ' + r'^Chapter\s+\d+', # 'Chapter ' followed by numbers + r'^\d+\.\s+', # Numbered chapters: '1. ', '2. ', etc. + r'^[A-Z\s]+$' # All caps headings + ] + + chapter_positions = [] + used_pattern = None + + for pattern in chapter_patterns: + if pattern is None: + continue + chapter_regex = re.compile(pattern, re.MULTILINE | re.IGNORECASE) + chapter_positions = [match.start() for match in chapter_regex.finditer(text)] + if chapter_positions: + used_pattern = pattern + break + + # If no chapters found, return the entire content as one chunk + if not chapter_positions: + metadata = get_chunk_metadata( + chunk=text, + full_text=text, + chunk_type="whole_document", + language=chunk_options.get('language', 'english') + ) + return [{'text': text, 'metadata': metadata}] + + # Split content into chapters + chunks = [] + for i in range(len(chapter_positions)): + start = chapter_positions[i] + end = chapter_positions[i + 1] if i + 1 < len(chapter_positions) else None + chapter = text[start:end] + + # Apply overlap if specified + if overlap > 0 and i > 0: + overlap_start = max(0, chapter_positions[i] - overlap) + chapter = text[overlap_start:end] + + chunks.append(chapter) + + # Post-process chunks + processed_chunks = post_process_chunks(chunks) + + # Add metadata to chunks + chunks_with_metadata = [] + for i, chunk in enumerate(processed_chunks): + metadata = get_chunk_metadata( + chunk=chunk, + full_text=text, + chunk_type="chapter", + chapter_number=i + 1, + chapter_pattern=used_pattern, + language=chunk_options.get('language', 'english') + ) + chunks_with_metadata.append({'text': chunk, 'metadata': metadata}) + + return chunks_with_metadata + +# +# End of ebook chapter chunking +####################################################################################################################### + +####################################################################################################################### +# +# Functions for adapative chunking: + +# FIXME - punkt + +def adaptive_chunk_size(text: str, base_size: int = 1000, min_size: int = 500, max_size: int = 2000) -> int: + # Tokenize the text into sentences + sentences = sent_tokenize(text) + + if not sentences: + return base_size + + # Calculate average sentence length + avg_sentence_length = sum(len(s.split()) for s in sentences) / len(sentences) + + # Adjust chunk size based on average sentence length + if avg_sentence_length < 10: + size_factor = 1.2 # Increase chunk size for short sentences + elif avg_sentence_length > 20: + size_factor = 0.8 # Decrease chunk size for long sentences + else: + size_factor = 1.0 + + # Calculate adaptive chunk size + adaptive_size = int(base_size * size_factor) + + # Ensure chunk size is within bounds + return max(min_size, min(adaptive_size, max_size)) + + +def adaptive_chunk_size_non_punkt(text: str, base_size: int, min_size: int = 100, max_size: int = 2000) -> int: + # Adaptive logic: adjust chunk size based on text complexity + words = text.split() + if not words: + return base_size # Return base_size if text is empty + + avg_word_length = sum(len(word) for word in words) / len(words) + + if avg_word_length > 6: # Threshold for "complex" text + adjusted_size = int(base_size * 0.8) # Reduce chunk size for complex text + elif avg_word_length < 4: # Threshold for "simple" text + adjusted_size = int(base_size * 1.2) # Increase chunk size for simple text + else: + adjusted_size = base_size + + # Ensure the chunk size is within the specified range + return max(min_size, min(adjusted_size, max_size)) + + +def adaptive_chunking(text: str, base_size: int = 1000, min_size: int = 500, max_size: int = 2000) -> List[str]: + logging.debug("adaptive_chunking...") + chunk_size = adaptive_chunk_size(text, base_size, min_size, max_size) + words = text.split() + chunks = [] + current_chunk = [] + current_length = 0 + + for word in words: + if current_length + len(word) > chunk_size and current_chunk: + chunks.append(' '.join(current_chunk)) + current_chunk = [] + current_length = 0 + current_chunk.append(word) + current_length += len(word) + 1 # +1 for space + + if current_chunk: + chunks.append(' '.join(current_chunk)) + + return chunks + +# FIXME - usage example +# chunk_options = { +# 'method': 'words', # or any other method +# 'base_size': 1000, +# 'min_size': 100, +# 'max_size': 2000, +# 'adaptive': True, +# 'language': 'en' +# } +#chunks = improved_chunking_process(your_text, chunk_options) + + +# Example of chunking a document with metadata +# document_metadata = { +# 'title': 'Example Document', +# 'author': 'John Doe', +# 'creation_date': '2023-06-14', +# 'source': 'https://example.com/document', +# 'document_type': 'article' +# } +# +# chunk_options = { +# 'method': 'sentences', +# 'base_size': 1000, +# 'adaptive': True, +# 'language': 'en' +# } +# +# processed_document = process_document_with_metadata(your_text, chunk_options, document_metadata) + + +# +# End of Chunking Library +####################################################################################################################### \ No newline at end of file diff --git a/App_Function_Libraries/DB/Character_Chat_DB.py b/App_Function_Libraries/DB/Character_Chat_DB.py new file mode 100644 index 0000000000000000000000000000000000000000..45f3376cbe4d4cca7ca00876c4507f447f9b35f7 --- /dev/null +++ b/App_Function_Libraries/DB/Character_Chat_DB.py @@ -0,0 +1,701 @@ +# character_chat_db.py +# Database functions for managing character cards and chat histories. +# # +# Imports +import configparser +import sqlite3 +import json +import os +import sys +from typing import List, Dict, Optional, Tuple, Any, Union + +from App_Function_Libraries.Utils.Utils import get_database_dir, get_project_relative_path, get_database_path +from Tests.Chat_APIs.Chat_APIs_Integration_test import logging + +# +####################################################################################################################### +# +# + +def ensure_database_directory(): + os.makedirs(get_database_dir(), exist_ok=True) + +ensure_database_directory() + + +# Construct the path to the config file +config_path = get_project_relative_path('Config_Files/config.txt') + +# Read the config file +config = configparser.ConfigParser() +config.read(config_path) + +# Get the chat db path from the config, or use the default if not specified +chat_DB_PATH = config.get('Database', 'chatDB_path', fallback=get_database_path('chatDB.db')) +print(f"Chat Database path: {chat_DB_PATH}") + +######################################################################################################## +# +# Functions + +# FIXME - Setup properly and test/add documentation for its existence... +def initialize_database(): + """Initialize the SQLite database with required tables and FTS5 virtual tables.""" + conn = None + try: + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + + # Enable foreign key constraints + cursor.execute("PRAGMA foreign_keys = ON;") + + # Create CharacterCards table with V2 fields + cursor.execute(""" + CREATE TABLE IF NOT EXISTS CharacterCards ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT UNIQUE NOT NULL, + description TEXT, + personality TEXT, + scenario TEXT, + image BLOB, + post_history_instructions TEXT, + first_mes TEXT, + mes_example TEXT, + creator_notes TEXT, + system_prompt TEXT, + alternate_greetings TEXT, + tags TEXT, + creator TEXT, + character_version TEXT, + extensions TEXT, + created_at DATETIME DEFAULT CURRENT_TIMESTAMP + ); + """) + + # Create CharacterChats table + cursor.execute(""" + CREATE TABLE IF NOT EXISTS CharacterChats ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + character_id INTEGER NOT NULL, + conversation_name TEXT, + chat_history TEXT, + is_snapshot BOOLEAN DEFAULT FALSE, + created_at DATETIME DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (character_id) REFERENCES CharacterCards(id) ON DELETE CASCADE + ); + """) + + # Create FTS5 virtual table for CharacterChats + cursor.execute(""" + CREATE VIRTUAL TABLE IF NOT EXISTS CharacterChats_fts USING fts5( + conversation_name, + chat_history, + content='CharacterChats', + content_rowid='id' + ); + """) + + # Create triggers to keep FTS5 table in sync with CharacterChats + cursor.executescript(""" + CREATE TRIGGER IF NOT EXISTS CharacterChats_ai AFTER INSERT ON CharacterChats BEGIN + INSERT INTO CharacterChats_fts(rowid, conversation_name, chat_history) + VALUES (new.id, new.conversation_name, new.chat_history); + END; + + CREATE TRIGGER IF NOT EXISTS CharacterChats_ad AFTER DELETE ON CharacterChats BEGIN + DELETE FROM CharacterChats_fts WHERE rowid = old.id; + END; + + CREATE TRIGGER IF NOT EXISTS CharacterChats_au AFTER UPDATE ON CharacterChats BEGIN + UPDATE CharacterChats_fts SET conversation_name = new.conversation_name, chat_history = new.chat_history + WHERE rowid = new.id; + END; + """) + + # Create ChatKeywords table + cursor.execute(""" + CREATE TABLE IF NOT EXISTS ChatKeywords ( + chat_id INTEGER NOT NULL, + keyword TEXT NOT NULL, + FOREIGN KEY (chat_id) REFERENCES CharacterChats(id) ON DELETE CASCADE + ); + """) + + # Create indexes for faster searches + cursor.execute(""" + CREATE INDEX IF NOT EXISTS idx_chatkeywords_keyword ON ChatKeywords(keyword); + """) + cursor.execute(""" + CREATE INDEX IF NOT EXISTS idx_chatkeywords_chat_id ON ChatKeywords(chat_id); + """) + + conn.commit() + logging.info("Database initialized successfully.") + except sqlite3.Error as e: + logging.error(f"SQLite error occurred during database initialization: {e}") + if conn: + conn.rollback() + raise + except Exception as e: + logging.error(f"Unexpected error occurred during database initialization: {e}") + if conn: + conn.rollback() + raise + finally: + if conn: + conn.close() + +# Call initialize_database() at the start of your application +def setup_chat_database(): + try: + initialize_database() + except Exception as e: + logging.critical(f"Failed to initialize database: {e}") + sys.exit(1) + +setup_chat_database() + +######################################################################################################## +# +# Character Card handling + +def parse_character_card(card_data: Dict[str, Any]) -> Dict[str, Any]: + """Parse and validate a character card according to V2 specification.""" + v2_data = { + 'name': card_data.get('name', ''), + 'description': card_data.get('description', ''), + 'personality': card_data.get('personality', ''), + 'scenario': card_data.get('scenario', ''), + 'first_mes': card_data.get('first_mes', ''), + 'mes_example': card_data.get('mes_example', ''), + 'creator_notes': card_data.get('creator_notes', ''), + 'system_prompt': card_data.get('system_prompt', ''), + 'post_history_instructions': card_data.get('post_history_instructions', ''), + 'alternate_greetings': json.dumps(card_data.get('alternate_greetings', [])), + 'tags': json.dumps(card_data.get('tags', [])), + 'creator': card_data.get('creator', ''), + 'character_version': card_data.get('character_version', ''), + 'extensions': json.dumps(card_data.get('extensions', {})) + } + + # Handle 'image' separately as it might be binary data + if 'image' in card_data: + v2_data['image'] = card_data['image'] + + return v2_data + + +def add_character_card(card_data: Dict[str, Any]) -> Optional[int]: + """Add or update a character card in the database.""" + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + try: + parsed_card = parse_character_card(card_data) + + # Check if character already exists + cursor.execute("SELECT id FROM CharacterCards WHERE name = ?", (parsed_card['name'],)) + row = cursor.fetchone() + + if row: + # Update existing character + character_id = row[0] + update_query = """ + UPDATE CharacterCards + SET description = ?, personality = ?, scenario = ?, image = ?, + post_history_instructions = ?, first_mes = ?, mes_example = ?, + creator_notes = ?, system_prompt = ?, alternate_greetings = ?, + tags = ?, creator = ?, character_version = ?, extensions = ? + WHERE id = ? + """ + cursor.execute(update_query, ( + parsed_card['description'], parsed_card['personality'], parsed_card['scenario'], + parsed_card['image'], parsed_card['post_history_instructions'], parsed_card['first_mes'], + parsed_card['mes_example'], parsed_card['creator_notes'], parsed_card['system_prompt'], + parsed_card['alternate_greetings'], parsed_card['tags'], parsed_card['creator'], + parsed_card['character_version'], parsed_card['extensions'], character_id + )) + else: + # Insert new character + insert_query = """ + INSERT INTO CharacterCards (name, description, personality, scenario, image, + post_history_instructions, first_mes, mes_example, creator_notes, system_prompt, + alternate_greetings, tags, creator, character_version, extensions) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """ + cursor.execute(insert_query, ( + parsed_card['name'], parsed_card['description'], parsed_card['personality'], + parsed_card['scenario'], parsed_card['image'], parsed_card['post_history_instructions'], + parsed_card['first_mes'], parsed_card['mes_example'], parsed_card['creator_notes'], + parsed_card['system_prompt'], parsed_card['alternate_greetings'], parsed_card['tags'], + parsed_card['creator'], parsed_card['character_version'], parsed_card['extensions'] + )) + character_id = cursor.lastrowid + + conn.commit() + return character_id + except sqlite3.IntegrityError as e: + logging.error(f"Error adding character card: {e}") + return None + except Exception as e: + logging.error(f"Unexpected error adding character card: {e}") + return None + finally: + conn.close() + +# def add_character_card(card_data: Dict) -> Optional[int]: +# """Add or update a character card in the database. +# +# Returns the ID of the inserted character or None if failed. +# """ +# conn = sqlite3.connect(chat_DB_PATH) +# cursor = conn.cursor() +# try: +# # Ensure all required fields are present +# required_fields = ['name', 'description', 'personality', 'scenario', 'image', 'post_history_instructions', 'first_message'] +# for field in required_fields: +# if field not in card_data: +# card_data[field] = '' # Assign empty string if field is missing +# +# # Check if character already exists +# cursor.execute("SELECT id FROM CharacterCards WHERE name = ?", (card_data['name'],)) +# row = cursor.fetchone() +# +# if row: +# # Update existing character +# character_id = row[0] +# cursor.execute(""" +# UPDATE CharacterCards +# SET description = ?, personality = ?, scenario = ?, image = ?, post_history_instructions = ?, first_message = ? +# WHERE id = ? +# """, ( +# card_data['description'], +# card_data['personality'], +# card_data['scenario'], +# card_data['image'], +# card_data['post_history_instructions'], +# card_data['first_message'], +# character_id +# )) +# else: +# # Insert new character +# cursor.execute(""" +# INSERT INTO CharacterCards (name, description, personality, scenario, image, post_history_instructions, first_message) +# VALUES (?, ?, ?, ?, ?, ?, ?) +# """, ( +# card_data['name'], +# card_data['description'], +# card_data['personality'], +# card_data['scenario'], +# card_data['image'], +# card_data['post_history_instructions'], +# card_data['first_message'] +# )) +# character_id = cursor.lastrowid +# +# conn.commit() +# return cursor.lastrowid +# except sqlite3.IntegrityError as e: +# logging.error(f"Error adding character card: {e}") +# return None +# except Exception as e: +# logging.error(f"Unexpected error adding character card: {e}") +# return None +# finally: +# conn.close() + + +def get_character_cards() -> List[Dict]: + """Retrieve all character cards from the database.""" + logging.debug(f"Fetching characters from DB: {chat_DB_PATH}") + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + cursor.execute("SELECT * FROM CharacterCards") + rows = cursor.fetchall() + columns = [description[0] for description in cursor.description] + conn.close() + characters = [dict(zip(columns, row)) for row in rows] + #logging.debug(f"Characters fetched from DB: {characters}") + return characters + + +def get_character_card_by_id(character_id: Union[int, Dict[str, Any]]) -> Optional[Dict[str, Any]]: + """ + Retrieve a single character card by its ID. + + Args: + character_id: Can be either an integer ID or a dictionary containing character data. + + Returns: + A dictionary containing the character card data, or None if not found. + """ + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + try: + if isinstance(character_id, dict): + # If a dictionary is passed, assume it's already a character card + return character_id + elif isinstance(character_id, int): + # If an integer is passed, fetch the character from the database + cursor.execute("SELECT * FROM CharacterCards WHERE id = ?", (character_id,)) + row = cursor.fetchone() + if row: + columns = [description[0] for description in cursor.description] + return dict(zip(columns, row)) + else: + logging.warning(f"Invalid type for character_id: {type(character_id)}") + return None + except Exception as e: + logging.error(f"Error in get_character_card_by_id: {e}") + return None + finally: + conn.close() + + +def update_character_card(character_id: int, card_data: Dict) -> bool: + """Update an existing character card.""" + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + try: + cursor.execute(""" + UPDATE CharacterCards + SET name = ?, description = ?, personality = ?, scenario = ?, image = ?, post_history_instructions = ?, first_message = ? + WHERE id = ? + """, ( + card_data.get('name'), + card_data.get('description'), + card_data.get('personality'), + card_data.get('scenario'), + card_data.get('image'), + card_data.get('post_history_instructions', ''), + card_data.get('first_message', "Hello! I'm ready to chat."), + character_id + )) + conn.commit() + return cursor.rowcount > 0 + except sqlite3.IntegrityError as e: + logging.error(f"Error updating character card: {e}") + return False + finally: + conn.close() + + +def delete_character_card(character_id: int) -> bool: + """Delete a character card and its associated chats.""" + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + try: + # Delete associated chats first due to foreign key constraint + cursor.execute("DELETE FROM CharacterChats WHERE character_id = ?", (character_id,)) + cursor.execute("DELETE FROM CharacterCards WHERE id = ?", (character_id,)) + conn.commit() + return cursor.rowcount > 0 + except sqlite3.Error as e: + logging.error(f"Error deleting character card: {e}") + return False + finally: + conn.close() + + +def add_character_chat(character_id: int, conversation_name: str, chat_history: List[Tuple[str, str]], keywords: Optional[List[str]] = None, is_snapshot: bool = False) -> Optional[int]: + """ + Add a new chat history for a character, optionally associating keywords. + + Args: + character_id (int): The ID of the character. + conversation_name (str): Name of the conversation. + chat_history (List[Tuple[str, str]]): List of (user, bot) message tuples. + keywords (Optional[List[str]]): List of keywords to associate with this chat. + is_snapshot (bool, optional): Whether this chat is a snapshot. + + Returns: + Optional[int]: The ID of the inserted chat or None if failed. + """ + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + try: + chat_history_json = json.dumps(chat_history) + cursor.execute(""" + INSERT INTO CharacterChats (character_id, conversation_name, chat_history, is_snapshot) + VALUES (?, ?, ?, ?) + """, ( + character_id, + conversation_name, + chat_history_json, + is_snapshot + )) + chat_id = cursor.lastrowid + + if keywords: + # Insert keywords into ChatKeywords table + keyword_records = [(chat_id, keyword.strip().lower()) for keyword in keywords] + cursor.executemany(""" + INSERT INTO ChatKeywords (chat_id, keyword) + VALUES (?, ?) + """, keyword_records) + + conn.commit() + return chat_id + except sqlite3.Error as e: + logging.error(f"Error adding character chat: {e}") + return None + finally: + conn.close() + + +def get_character_chats(character_id: Optional[int] = None) -> List[Dict]: + """Retrieve all chats, or chats for a specific character if character_id is provided.""" + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + if character_id is not None: + cursor.execute("SELECT * FROM CharacterChats WHERE character_id = ?", (character_id,)) + else: + cursor.execute("SELECT * FROM CharacterChats") + rows = cursor.fetchall() + columns = [description[0] for description in cursor.description] + conn.close() + return [dict(zip(columns, row)) for row in rows] + + +def get_character_chat_by_id(chat_id: int) -> Optional[Dict]: + """Retrieve a single chat by its ID.""" + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + cursor.execute("SELECT * FROM CharacterChats WHERE id = ?", (chat_id,)) + row = cursor.fetchone() + conn.close() + if row: + columns = [description[0] for description in cursor.description] + chat = dict(zip(columns, row)) + chat['chat_history'] = json.loads(chat['chat_history']) + return chat + return None + + +def search_character_chats(query: str, character_id: Optional[int] = None) -> Tuple[List[Dict], str]: + """ + Search for character chats using FTS5, optionally filtered by character_id. + + Args: + query (str): The search query. + character_id (Optional[int]): The ID of the character to filter chats by. + + Returns: + Tuple[List[Dict], str]: A list of matching chats and a status message. + """ + if not query.strip(): + return [], "Please enter a search query." + + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + try: + if character_id is not None: + # Search with character_id filter + cursor.execute(""" + SELECT CharacterChats.id, CharacterChats.conversation_name, CharacterChats.chat_history + FROM CharacterChats_fts + JOIN CharacterChats ON CharacterChats_fts.rowid = CharacterChats.id + WHERE CharacterChats_fts MATCH ? AND CharacterChats.character_id = ? + ORDER BY rank + """, (query, character_id)) + else: + # Search without character_id filter + cursor.execute(""" + SELECT CharacterChats.id, CharacterChats.conversation_name, CharacterChats.chat_history + FROM CharacterChats_fts + JOIN CharacterChats ON CharacterChats_fts.rowid = CharacterChats.id + WHERE CharacterChats_fts MATCH ? + ORDER BY rank + """, (query,)) + + rows = cursor.fetchall() + columns = [description[0] for description in cursor.description] + results = [dict(zip(columns, row)) for row in rows] + + if character_id is not None: + status_message = f"Found {len(results)} chat(s) matching '{query}' for the selected character." + else: + status_message = f"Found {len(results)} chat(s) matching '{query}' across all characters." + + return results, status_message + except Exception as e: + logging.error(f"Error searching chats with FTS5: {e}") + return [], f"Error occurred during search: {e}" + finally: + conn.close() + +def update_character_chat(chat_id: int, chat_history: List[Tuple[str, str]]) -> bool: + """Update an existing chat history.""" + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + try: + chat_history_json = json.dumps(chat_history) + cursor.execute(""" + UPDATE CharacterChats + SET chat_history = ? + WHERE id = ? + """, ( + chat_history_json, + chat_id + )) + conn.commit() + return cursor.rowcount > 0 + except sqlite3.Error as e: + logging.error(f"Error updating character chat: {e}") + return False + finally: + conn.close() + + +def delete_character_chat(chat_id: int) -> bool: + """Delete a specific chat.""" + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + try: + cursor.execute("DELETE FROM CharacterChats WHERE id = ?", (chat_id,)) + conn.commit() + return cursor.rowcount > 0 + except sqlite3.Error as e: + logging.error(f"Error deleting character chat: {e}") + return False + finally: + conn.close() + +def fetch_keywords_for_chats(keywords: List[str]) -> List[int]: + """ + Fetch chat IDs associated with any of the specified keywords. + + Args: + keywords (List[str]): List of keywords to search for. + + Returns: + List[int]: List of chat IDs associated with the keywords. + """ + if not keywords: + return [] + + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + try: + # Construct the WHERE clause to search for each keyword + keyword_clauses = " OR ".join(["keyword = ?"] * len(keywords)) + sql_query = f"SELECT DISTINCT chat_id FROM ChatKeywords WHERE {keyword_clauses}" + cursor.execute(sql_query, keywords) + rows = cursor.fetchall() + chat_ids = [row[0] for row in rows] + return chat_ids + except Exception as e: + logging.error(f"Error in fetch_keywords_for_chats: {e}") + return [] + finally: + conn.close() + +def save_chat_history_to_character_db(character_id: int, conversation_name: str, chat_history: List[Tuple[str, str]]) -> Optional[int]: + """Save chat history to the CharacterChats table. + + Returns the ID of the inserted chat or None if failed. + """ + return add_character_chat(character_id, conversation_name, chat_history) + +def migrate_chat_to_media_db(): + pass + + +def search_db(query: str, fields: List[str], where_clause: str = "", page: int = 1, results_per_page: int = 5) -> List[Dict[str, Any]]: + """ + Perform a full-text search on specified fields with optional filtering and pagination. + + Args: + query (str): The search query. + fields (List[str]): List of fields to search in. + where_clause (str, optional): Additional SQL WHERE clause to filter results. + page (int, optional): Page number for pagination. + results_per_page (int, optional): Number of results per page. + + Returns: + List[Dict[str, Any]]: List of matching chat records with content and metadata. + """ + if not query.strip(): + return [] + + conn = sqlite3.connect(chat_DB_PATH) + cursor = conn.cursor() + try: + # Construct the MATCH query for FTS5 + match_query = " AND ".join(fields) + f" MATCH ?" + # Adjust the query with the fields + fts_query = f""" + SELECT CharacterChats.id, CharacterChats.conversation_name, CharacterChats.chat_history + FROM CharacterChats_fts + JOIN CharacterChats ON CharacterChats_fts.rowid = CharacterChats.id + WHERE {match_query} + """ + if where_clause: + fts_query += f" AND ({where_clause})" + fts_query += " ORDER BY rank LIMIT ? OFFSET ?" + offset = (page - 1) * results_per_page + cursor.execute(fts_query, (query, results_per_page, offset)) + rows = cursor.fetchall() + columns = [description[0] for description in cursor.description] + results = [dict(zip(columns, row)) for row in rows] + return results + except Exception as e: + logging.error(f"Error in search_db: {e}") + return [] + finally: + conn.close() + + +def perform_full_text_search_chat(query: str, relevant_chat_ids: List[int], page: int = 1, results_per_page: int = 5) -> \ +List[Dict[str, Any]]: + """ + Perform a full-text search within the specified chat IDs using FTS5. + + Args: + query (str): The user's query. + relevant_chat_ids (List[int]): List of chat IDs to search within. + page (int): Pagination page number. + results_per_page (int): Number of results per page. + + Returns: + List[Dict[str, Any]]: List of search results with content and metadata. + """ + try: + # Construct a WHERE clause to limit the search to relevant chat IDs + where_clause = " OR ".join([f"media_id = {chat_id}" for chat_id in relevant_chat_ids]) + if not where_clause: + where_clause = "1" # No restriction if no chat IDs + + # Perform full-text search using FTS5 + fts_results = search_db(query, ["content"], where_clause, page=page, results_per_page=results_per_page) + + filtered_fts_results = [ + { + "content": result['content'], + "metadata": {"media_id": result['id']} + } + for result in fts_results + if result['id'] in relevant_chat_ids + ] + return filtered_fts_results + except Exception as e: + logging.error(f"Error in perform_full_text_search_chat: {str(e)}") + return [] + + +def fetch_all_chats() -> List[Dict[str, Any]]: + """ + Fetch all chat messages from the database. + + Returns: + List[Dict[str, Any]]: List of chat messages with relevant metadata. + """ + try: + chats = get_character_chats() # Modify this function to retrieve all chats + return chats + except Exception as e: + logging.error(f"Error fetching all chats: {str(e)}") + return [] + +# +# End of Character_Chat_DB.py +####################################################################################################################### diff --git a/App_Function_Libraries/DB/DB_Manager.py b/App_Function_Libraries/DB/DB_Manager.py new file mode 100644 index 0000000000000000000000000000000000000000..a11e4d9a3872d1f8ba36e707ca4d5914332c7675 --- /dev/null +++ b/App_Function_Libraries/DB/DB_Manager.py @@ -0,0 +1,991 @@ +# DB_Manager.py +# Description: This file contains the DatabaseManager class, which is responsible for managing the database connection, i.e. either SQLite or Elasticsearch. +# +# Imports +import configparser +import os +import logging +import time +from typing import Tuple, List, Union, Dict +# +# 3rd-Party Libraries +from elasticsearch import Elasticsearch +# +# Import your existing SQLite functions +from App_Function_Libraries.DB.SQLite_DB import DatabaseError +from App_Function_Libraries.DB.SQLite_DB import ( + update_media_content as sqlite_update_media_content, + list_prompts as sqlite_list_prompts, + search_and_display as sqlite_search_and_display, + fetch_prompt_details as sqlite_fetch_prompt_details, + keywords_browser_interface as sqlite_keywords_browser_interface, + add_keyword as sqlite_add_keyword, + delete_keyword as sqlite_delete_keyword, + export_keywords_to_csv as sqlite_export_keywords_to_csv, + ingest_article_to_db as sqlite_ingest_article_to_db, + add_media_to_database as sqlite_add_media_to_database, + import_obsidian_note_to_db as sqlite_import_obsidian_note_to_db, + add_prompt as sqlite_add_prompt, + delete_chat_message as sqlite_delete_chat_message, + update_chat_message as sqlite_update_chat_message, + add_chat_message as sqlite_add_chat_message, + get_chat_messages as sqlite_get_chat_messages, + search_chat_conversations as sqlite_search_chat_conversations, + create_chat_conversation as sqlite_create_chat_conversation, + save_chat_history_to_database as sqlite_save_chat_history_to_database, + view_database as sqlite_view_database, + get_transcripts as sqlite_get_transcripts, + get_trashed_items as sqlite_get_trashed_items, + user_delete_item as sqlite_user_delete_item, + empty_trash as sqlite_empty_trash, + create_automated_backup as sqlite_create_automated_backup, + add_or_update_prompt as sqlite_add_or_update_prompt, + load_prompt_details as sqlite_load_prompt_details, + load_preset_prompts as sqlite_load_preset_prompts, + insert_prompt_to_db as sqlite_insert_prompt_to_db, + delete_prompt as sqlite_delete_prompt, + search_and_display_items as sqlite_search_and_display_items, + get_conversation_name as sqlite_get_conversation_name, + add_media_with_keywords as sqlite_add_media_with_keywords, + check_media_and_whisper_model as sqlite_check_media_and_whisper_model, \ + create_document_version as sqlite_create_document_version, + get_document_version as sqlite_get_document_version, sqlite_search_db, add_media_chunk as sqlite_add_media_chunk, + sqlite_update_fts_for_media, get_unprocessed_media as sqlite_get_unprocessed_media, fetch_item_details as sqlite_fetch_item_details, \ + search_media_database as sqlite_search_media_database, mark_as_trash as sqlite_mark_as_trash, \ + get_media_transcripts as sqlite_get_media_transcripts, get_specific_transcript as sqlite_get_specific_transcript, \ + get_media_summaries as sqlite_get_media_summaries, get_specific_summary as sqlite_get_specific_summary, \ + get_media_prompts as sqlite_get_media_prompts, get_specific_prompt as sqlite_get_specific_prompt, \ + delete_specific_transcript as sqlite_delete_specific_transcript, + delete_specific_summary as sqlite_delete_specific_summary, \ + delete_specific_prompt as sqlite_delete_specific_prompt, + fetch_keywords_for_media as sqlite_fetch_keywords_for_media, \ + update_keywords_for_media as sqlite_update_keywords_for_media, check_media_exists as sqlite_check_media_exists, \ + search_prompts as sqlite_search_prompts, get_media_content as sqlite_get_media_content, \ + get_paginated_files as sqlite_get_paginated_files, get_media_title as sqlite_get_media_title, \ + get_all_content_from_database as sqlite_get_all_content_from_database, + get_next_media_id as sqlite_get_next_media_id, \ + batch_insert_chunks as sqlite_batch_insert_chunks, Database, save_workflow_chat_to_db as sqlite_save_workflow_chat_to_db, \ + get_workflow_chat as sqlite_get_workflow_chat, update_media_content_with_version as sqlite_update_media_content_with_version, \ + check_existing_media as sqlite_check_existing_media, get_all_document_versions as sqlite_get_all_document_versions, \ + fetch_paginated_data as sqlite_fetch_paginated_data, get_latest_transcription as sqlite_get_latest_transcription, \ + mark_media_as_processed as sqlite_mark_media_as_processed, +) +from App_Function_Libraries.DB.Character_Chat_DB import ( + add_character_card as sqlite_add_character_card, get_character_cards as sqlite_get_character_cards, \ + get_character_card_by_id as sqlite_get_character_card_by_id, update_character_card as sqlite_update_character_card, \ + delete_character_card as sqlite_delete_character_card, add_character_chat as sqlite_add_character_chat, \ + get_character_chats as sqlite_get_character_chats, get_character_chat_by_id as sqlite_get_character_chat_by_id, \ + update_character_chat as sqlite_update_character_chat, delete_character_chat as sqlite_delete_character_chat, \ + migrate_chat_to_media_db as sqlite_migrate_chat_to_media_db, +) +# +# Local Imports +from App_Function_Libraries.Utils.Utils import load_comprehensive_config, get_database_path, get_project_relative_path +# +# End of imports +############################################################################################################ + + +############################################################################################################ +# +# Database Config loading + +logger = logging.getLogger(__name__) + +config_path = get_project_relative_path('Config_Files/config.txt') +config = configparser.ConfigParser() +config.read(config_path) + +db_path: str = config.get('Database', 'sqlite_path', fallback='./Databases/media_summary.db') +backup_path: str = config.get('Database', 'backup_path', fallback='database_backups') +backup_dir: Union[str, bytes] = os.environ.get('DB_BACKUP_DIR', backup_path) + +def get_db_config(): + try: + config = load_comprehensive_config() + + if 'Database' not in config: + print("Warning: 'Database' section not found in config. Using default values.") + return default_db_config() + + return { + 'type': config.get('Database', 'type', fallback='sqlite'), + 'sqlite_path': config.get('Database', 'sqlite_path', fallback='Databases/media_summary.db'), + 'elasticsearch_host': config.get('Database', 'elasticsearch_host', fallback='localhost'), + 'elasticsearch_port': config.getint('Database', 'elasticsearch_port', fallback=9200) + } + except FileNotFoundError: + print("Warning: Config file not found. Using default database configuration.") + return default_db_config() + except Exception as e: + print(f"Error reading config: {str(e)}. Using default database configuration.") + return default_db_config() + +def default_db_config(): + return { + 'type': 'sqlite', + 'sqlite_path': get_database_path('media_summary.db'), + 'elasticsearch_host': 'localhost', + 'elasticsearch_port': 9200 + } + +def ensure_directory_exists(file_path): + directory = os.path.dirname(file_path) + if not os.path.exists(directory): + os.makedirs(directory) + print(f"Created directory: {directory}") + +db_config = get_db_config() +db_type = db_config['type'] + +if db_type == 'sqlite': + db = Database(os.path.basename(db_config['sqlite_path'])) +elif db_type == 'elasticsearch': + raise NotImplementedError("Elasticsearch support not yet implemented") +else: + raise ValueError(f"Unsupported database type: {db_type}") + +print(f"Database path: {db.db_path}") + +def get_db_config(): + try: + config = load_comprehensive_config() + + if 'Database' not in config: + print("Warning: 'Database' section not found in config. Using default values.") + return default_db_config() + + return { + 'type': config.get('Database', 'type', fallback='sqlite'), + 'sqlite_path': config.get('Database', 'sqlite_path', fallback='Databases/media_summary.db'), + 'elasticsearch_host': config.get('Database', 'elasticsearch_host', fallback='localhost'), + 'elasticsearch_port': config.getint('Database', 'elasticsearch_port', fallback=9200) + } + except FileNotFoundError: + print("Warning: Config file not found. Using default database configuration.") + return default_db_config() + except Exception as e: + print(f"Error reading config: {str(e)}. Using default database configuration.") + return default_db_config() + + +def default_db_config(): + """Return the default database configuration with project-relative paths.""" + return { + 'type': 'sqlite', + 'sqlite_path': get_database_path('media_summary.db'), + 'elasticsearch_host': 'localhost', + 'elasticsearch_port': 9200 + } + + +def ensure_directory_exists(file_path): + directory = os.path.dirname(file_path) + if not os.path.exists(directory): + os.makedirs(directory) + print(f"Created directory: {directory}") + +# Use the config to set up the database +db_config = get_db_config() +db_type = db_config['type'] + +if db_type == 'sqlite': + db = Database(os.path.basename(db_config['sqlite_path'])) +elif db_type == 'elasticsearch': + # Implement Elasticsearch setup here if needed + raise NotImplementedError("Elasticsearch support not yet implemented") +else: + raise ValueError(f"Unsupported database type: {db_type}") + +# Print database path for debugging +print(f"Database path: {db.db_path}") + +# Sanity Check for SQLite DB +# FIXME - Remove this after testing / Writing Unit tests +# try: +# db.execute_query("CREATE TABLE IF NOT EXISTS test_table (id INTEGER PRIMARY KEY)") +# logger.info("Successfully created test table") +# except DatabaseError as e: +# logger.error(f"Failed to create test table: {e}") + +# +# End of Database Config loading +############################################################################################################ +# +# DB Search functions + +def search_db(search_query: str, search_fields: List[str], keywords: str, page: int = 1, results_per_page: int = 10): + if db_type == 'sqlite': + return sqlite_search_db(search_query, search_fields, keywords, page, results_per_page) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version when available + raise NotImplementedError("Elasticsearch version of search_db not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +def view_database(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_view_database(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def search_and_display_items(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_search_and_display_items(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def get_all_content_from_database(): + if db_type == 'sqlite': + return sqlite_get_all_content_from_database() + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def search_and_display(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_search_and_display(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def check_media_exists(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_check_media_exists(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def get_paginated_files(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_paginated_files(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def get_media_title(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_media_title(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def get_next_media_id(): + if db_type == 'sqlite': + return sqlite_get_next_media_id() + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +# +# End of DB-Searching functions +############################################################################################################ + + +############################################################################################################ +# +# Transcript-related Functions + +def get_transcripts(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_transcripts(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +# +# End of Transcript-related Functions +############################################################################################################ + + +############################################################################################################ +# +# DB-Ingestion functions + +def add_media_to_database(*args, **kwargs): + if db_type == 'sqlite': + result = sqlite_add_media_to_database(*args, **kwargs) + + # Extract content + segments = kwargs.get('segments') if 'segments' in kwargs else args[2] if len(args) > 2 else None + if segments is None: + raise ValueError("Segments not provided in arguments") + + if isinstance(segments, list): + content = ' '.join([segment.get('Text', '') for segment in segments if 'Text' in segment]) + elif isinstance(segments, dict): + content = segments.get('text', '') or segments.get('content', '') + else: + content = str(segments) + + # Extract media_id from the result + # Assuming the result is in the format "Media 'Title' added/updated successfully with ID: {media_id}" + import re + match = re.search(r"with ID: (\d+)", result) + if match: + media_id = int(match.group(1)) + + # Create initial document version + sqlite_create_document_version(media_id, content) + + return result + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_to_database not yet implemented") + +def check_existing_media(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_check_existing_media(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of check_existing_media not yet implemented") + +def update_media_content_with_version(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_update_media_content_with_version(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of update_media_content not yet implemented") + +def import_obsidian_note_to_db(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_import_obsidian_note_to_db(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + + +def update_media_content(*args, **kwargs): + if db_type == 'sqlite': + result = sqlite_update_media_content(*args, **kwargs) + + # Extract media_id and content + selected_item = args[0] + item_mapping = args[1] + content_input = args[2] + + if selected_item and item_mapping and selected_item in item_mapping: + media_id = item_mapping[selected_item] + + # Create new document version + sqlite_create_document_version(media_id, content_input) + + return result + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of update_media_content not yet implemented") + + +def add_media_with_keywords(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_add_media_with_keywords(*args, **kwargs) + elif db_type == 'elasticsearch': + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def check_media_and_whisper_model(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_check_media_and_whisper_model(*args, **kwargs) + elif db_type == 'elasticsearch': + raise NotImplementedError("Elasticsearch version of check_media_and_whisper_model not yet implemented") + +def ingest_article_to_db(url, title, author, content, keywords, summary, ingestion_date, custom_prompt): + if db_type == 'sqlite': + return sqlite_ingest_article_to_db(url, title, author, content, keywords, summary, ingestion_date, custom_prompt) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of ingest_article_to_db not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + + +def add_media_chunk(*args, **kwargs): + if db_type == 'sqlite': + sqlite_add_media_chunk(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +def batch_insert_chunks(*args, **kwargs): + if db_type == 'sqlite': + sqlite_batch_insert_chunks(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +def update_fts_for_media(media_id: int): + if db_type == 'sqlite': + sqlite_update_fts_for_media(db, media_id) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + + +def get_unprocessed_media(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_unprocessed_media(db) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of get_unprocessed_media not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + + +def mark_media_as_processed(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_mark_media_as_processed(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of mark_media_as_processed not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + + +# +# End of DB-Ingestion functions +############################################################################################################ + + +############################################################################################################ +# +# Prompt-related functions #FIXME rename /resort + +def list_prompts(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_list_prompts(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def search_prompts(query): + if db_type == 'sqlite': + return sqlite_search_prompts(query) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def fetch_prompt_details(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_fetch_prompt_details(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def add_prompt(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_add_prompt(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + + +def add_or_update_prompt(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_add_or_update_prompt(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def load_prompt_details(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_load_prompt_details(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def load_preset_prompts(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_load_preset_prompts() + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def insert_prompt_to_db(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_insert_prompt_to_db(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def delete_prompt(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_delete_prompt(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def search_media_database(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_search_media_database(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version when available + raise NotImplementedError("Elasticsearch version of search_media_database not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +def mark_as_trash(media_id: int) -> None: + if db_type == 'sqlite': + return sqlite_mark_as_trash(media_id) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version when available + raise NotImplementedError("Elasticsearch version of mark_as_trash not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + + +def get_latest_transcription(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_latest_transcription(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of get_latest_transcription not yet implemented") + +def fetch_paginated_data(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_fetch_paginated_data(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of fetch_paginated_data not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + + +def get_media_content(media_id: int) -> str: + if db_type == 'sqlite': + return sqlite_get_media_content(media_id) + elif db_type == 'elasticsearch': + raise NotImplementedError("Elasticsearch version of get_media_content not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +def get_media_transcripts(media_id: int) -> List[Dict]: + if db_type == 'sqlite': + return sqlite_get_media_transcripts(media_id) + elif db_type == 'elasticsearch': + raise NotImplementedError("Elasticsearch version of get_media_transcripts not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +def get_specific_transcript(transcript_id: int) -> Dict: + if db_type == 'sqlite': + return sqlite_get_specific_transcript(transcript_id) + elif db_type == 'elasticsearch': + raise NotImplementedError("Elasticsearch version of get_specific_transcript not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +def get_media_summaries(media_id: int) -> List[Dict]: + if db_type == 'sqlite': + return sqlite_get_media_summaries(media_id) + elif db_type == 'elasticsearch': + raise NotImplementedError("Elasticsearch version of get_media_summaries not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +def get_specific_summary(summary_id: int) -> Dict: + if db_type == 'sqlite': + return sqlite_get_specific_summary(summary_id) + elif db_type == 'elasticsearch': + raise NotImplementedError("Elasticsearch version of get_specific_summary not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +def fetch_item_details_single(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_fetch_item_details(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of fetch_item_details not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +def get_all_document_versions(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_all_document_versions(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of get_all_document_versions not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") +# +# +############################################################################################################ +# +# Prompt Functions: + +def get_media_prompts(media_id: int) -> List[Dict]: + if db_type == 'sqlite': + return sqlite_get_media_prompts(media_id) + elif db_type == 'elasticsearch': + raise NotImplementedError("Elasticsearch version of get_media_prompts not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +def get_specific_prompt(prompt_id: int) -> Dict: + if db_type == 'sqlite': + return sqlite_get_specific_prompt(prompt_id) + elif db_type == 'elasticsearch': + raise NotImplementedError("Elasticsearch version of get_specific_prompt not yet implemented") + else: + return {'error': f"Unsupported database type: {db_type}"} + +def delete_specific_transcript(transcript_id: int) -> str: + if db_type == 'sqlite': + return sqlite_delete_specific_transcript(transcript_id) + elif db_type == 'elasticsearch': + raise NotImplementedError("Elasticsearch version of delete_specific_transcript not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +def delete_specific_summary(summary_id: int) -> str: + if db_type == 'sqlite': + return sqlite_delete_specific_summary(summary_id) + elif db_type == 'elasticsearch': + raise NotImplementedError("Elasticsearch version of delete_specific_summary not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +def delete_specific_prompt(prompt_id: int) -> str: + if db_type == 'sqlite': + return sqlite_delete_specific_prompt(prompt_id) + elif db_type == 'elasticsearch': + raise NotImplementedError("Elasticsearch version of delete_specific_prompt not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + + +# +# End of Prompt-related functions +############################################################################################################ + +############################################################################################################ +# +# Keywords-related Functions + +def keywords_browser_interface(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_keywords_browser_interface() + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def add_keyword(*args, **kwargs): + if db_type == 'sqlite': + with db.get_connection() as conn: + cursor = conn.cursor() + return sqlite_add_keyword(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def delete_keyword(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_delete_keyword(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def export_keywords_to_csv(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_export_keywords_to_csv() + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def update_keywords_for_media(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_update_keywords_for_media(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def fetch_keywords_for_media(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_fetch_keywords_for_media(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +# +# End of Keywords-related Functions +############################################################################################################ + +############################################################################################################ +# +# Chat-related Functions + +def delete_chat_message(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_delete_chat_message(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def update_chat_message(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_update_chat_message(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def add_chat_message(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_add_chat_message(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def get_chat_messages(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_chat_messages(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def search_chat_conversations(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_search_chat_conversations(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def create_chat_conversation(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_create_chat_conversation(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def save_chat_history_to_database(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_save_chat_history_to_database(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def get_conversation_name(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_conversation_name(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +# +# End of Chat-related Functions +############################################################################################################ + + +############################################################################################################ +# +# Character Chat-related Functions + +def add_character_card(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_add_character_card(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_character_card not yet implemented") + +def get_character_cards(): + if db_type == 'sqlite': + return sqlite_get_character_cards() + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of get_character_cards not yet implemented") + +def get_character_card_by_id(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_character_card_by_id(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of get_character_card_by_id not yet implemented") + +def update_character_card(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_update_character_card(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of update_character_card not yet implemented") + +def delete_character_card(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_delete_character_card(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of delete_character_card not yet implemented") + +def add_character_chat(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_add_character_chat(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_character_chat not yet implemented") + +def get_character_chats(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_character_chats(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of get_character_chats not yet implemented") + +def get_character_chat_by_id(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_character_chat_by_id(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of get_character_chat_by_id not yet implemented") + +def update_character_chat(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_update_character_chat(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of update_character_chat not yet implemented") + +def delete_character_chat(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_delete_character_chat(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of delete_character_chat not yet implemented") + +def migrate_chat_to_media_db(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_migrate_chat_to_media_db(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of migrate_chat_to_media_db not yet implemented") + +# +# End of Character Chat-related Functions +############################################################################################################ + + +############################################################################################################ +# +# Trash-related Functions + +def get_trashed_items(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_trashed_items() + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def user_delete_item(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_user_delete_item(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +def empty_trash(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_empty_trash(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + + +def fetch_item_details(media_id: int) -> Tuple[str, str, str]: + """ + Fetch the details of a media item including content, prompt, and summary. + + Args: + media_id (int): The ID of the media item. + + Returns: + Tuple[str, str, str]: A tuple containing (content, prompt, summary). + If an error occurs, it returns empty strings for each field. + """ + if db_type == 'sqlite': + return sqlite_fetch_item_details(media_id) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version when available + raise NotImplementedError("Elasticsearch version of fetch_item_details not yet implemented") + else: + raise ValueError(f"Unsupported database type: {db_type}") + +# +# End of Trash-related Functions +############################################################################################################ + + +############################################################################################################ +# +# DB-Backup Functions + +def create_automated_backup(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_create_automated_backup(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented") + +# +# End of DB-Backup Functions +############################################################################################################ + + +############################################################################################################ +# +# Document Versioning Functions + +def create_document_version(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_create_document_version(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of create_document_version not yet implemented") + +def get_document_version(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_document_version(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of get_document_version not yet implemented") + +# +# End of Document Versioning Functions +############################################################################################################ + + +############################################################################################################ +# +# Workflow Functions + +def get_workflow_chat(*args, **kwargs): + if db_type == 'sqlite': + return sqlite_get_workflow_chat(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of get_workflow_chat not yet implemented") + + +def save_workflow_chat_to_db(*args, **kwargs): + if db_type == 'sqlite': + # FIXME + return sqlite_save_workflow_chat_to_db(*args, **kwargs) + elif db_type == 'elasticsearch': + # Implement Elasticsearch version + raise NotImplementedError("Elasticsearch version of save_workflow_chat_to_db not yet implemented") + +# +# End of Workflow Functions +############################################################################################################ + +# Dead code FIXME +# def close_connection(): +# if db_type == 'sqlite': +# db.get_connection().close() + +# +# End of file +############################################################################################################ diff --git a/App_Function_Libraries/DB/RAG_QA_Chat_DB.py b/App_Function_Libraries/DB/RAG_QA_Chat_DB.py new file mode 100644 index 0000000000000000000000000000000000000000..6622ac5980bea0731894c257640f442052eb66b3 --- /dev/null +++ b/App_Function_Libraries/DB/RAG_QA_Chat_DB.py @@ -0,0 +1,722 @@ +# RAG_QA_Chat_DB.py +# Description: This file contains the database operations for the RAG QA Chat + Notes system. +# +# Imports +import configparser +import logging +import re +import sqlite3 +import uuid +from contextlib import contextmanager +from datetime import datetime + +from App_Function_Libraries.Utils.Utils import get_project_relative_path, get_database_path + +# +# External Imports +# (No external imports) +# +# Local Imports +# (No additional local imports) +# +######################################################################################################################## +# +# Functions: + +# Construct the path to the config file +config_path = get_project_relative_path('Config_Files/config.txt') + +# Read the config file +config = configparser.ConfigParser() +config.read(config_path) + +# Get the SQLite path from the config, or use the default if not specified +if config.has_section('Database') and config.has_option('Database', 'rag_qa_db_path'): + rag_qa_db_path = config.get('Database', 'rag_qa_db_path') +else: + rag_qa_db_path = get_database_path('RAG_QA_Chat.db') + +print(f"RAG QA Chat Database path: {rag_qa_db_path}") + +# Set up logging +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + +# Database schema +SCHEMA_SQL = ''' +-- Table for storing chat messages +CREATE TABLE IF NOT EXISTS rag_qa_chats ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + conversation_id TEXT NOT NULL, + timestamp DATETIME NOT NULL, + role TEXT NOT NULL, + content TEXT NOT NULL +); + +-- Table for storing conversation metadata +CREATE TABLE IF NOT EXISTS conversation_metadata ( + conversation_id TEXT PRIMARY KEY, + created_at DATETIME NOT NULL, + last_updated DATETIME NOT NULL, + title TEXT NOT NULL +); + +-- Table for storing keywords +CREATE TABLE IF NOT EXISTS rag_qa_keywords ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + keyword TEXT NOT NULL UNIQUE +); + +-- Table for linking keywords to conversations +CREATE TABLE IF NOT EXISTS rag_qa_conversation_keywords ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + conversation_id TEXT NOT NULL, + keyword_id INTEGER NOT NULL, + FOREIGN KEY (conversation_id) REFERENCES conversation_metadata(conversation_id), + FOREIGN KEY (keyword_id) REFERENCES rag_qa_keywords(id) +); + +-- Table for storing keyword collections +CREATE TABLE IF NOT EXISTS rag_qa_keyword_collections ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT NOT NULL UNIQUE, + parent_id INTEGER, + FOREIGN KEY (parent_id) REFERENCES rag_qa_keyword_collections(id) +); + +-- Table for linking keywords to collections +CREATE TABLE IF NOT EXISTS rag_qa_collection_keywords ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + collection_id INTEGER NOT NULL, + keyword_id INTEGER NOT NULL, + FOREIGN KEY (collection_id) REFERENCES rag_qa_keyword_collections(id), + FOREIGN KEY (keyword_id) REFERENCES rag_qa_keywords(id) +); + +-- Table for storing notes +CREATE TABLE IF NOT EXISTS rag_qa_notes ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + conversation_id TEXT NOT NULL, + title TEXT NOT NULL, + content TEXT NOT NULL, + timestamp DATETIME NOT NULL, + FOREIGN KEY (conversation_id) REFERENCES conversation_metadata(conversation_id) +); + +-- Table for linking notes to keywords +CREATE TABLE IF NOT EXISTS rag_qa_note_keywords ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + note_id INTEGER NOT NULL, + keyword_id INTEGER NOT NULL, + FOREIGN KEY (note_id) REFERENCES rag_qa_notes(id), + FOREIGN KEY (keyword_id) REFERENCES rag_qa_keywords(id) +); + +-- Indexes for improved query performance +CREATE INDEX IF NOT EXISTS idx_rag_qa_chats_conversation_id ON rag_qa_chats(conversation_id); +CREATE INDEX IF NOT EXISTS idx_rag_qa_chats_timestamp ON rag_qa_chats(timestamp); +CREATE INDEX IF NOT EXISTS idx_rag_qa_keywords_keyword ON rag_qa_keywords(keyword); +CREATE INDEX IF NOT EXISTS idx_rag_qa_conversation_keywords_conversation_id ON rag_qa_conversation_keywords(conversation_id); +CREATE INDEX IF NOT EXISTS idx_rag_qa_conversation_keywords_keyword_id ON rag_qa_conversation_keywords(keyword_id); +CREATE INDEX IF NOT EXISTS idx_rag_qa_keyword_collections_parent_id ON rag_qa_keyword_collections(parent_id); +CREATE INDEX IF NOT EXISTS idx_rag_qa_collection_keywords_collection_id ON rag_qa_collection_keywords(collection_id); +CREATE INDEX IF NOT EXISTS idx_rag_qa_collection_keywords_keyword_id ON rag_qa_collection_keywords(keyword_id); + +-- Full-text search virtual table for chat content +CREATE VIRTUAL TABLE IF NOT EXISTS rag_qa_chats_fts USING fts5(conversation_id, timestamp, role, content); + +-- Trigger to keep the FTS table up to date +CREATE TRIGGER IF NOT EXISTS rag_qa_chats_ai AFTER INSERT ON rag_qa_chats BEGIN + INSERT INTO rag_qa_chats_fts(conversation_id, timestamp, role, content) VALUES (new.conversation_id, new.timestamp, new.role, new.content); +END; +''' + +# Database connection management +@contextmanager +def get_db_connection(): + conn = sqlite3.connect(rag_qa_db_path) + try: + yield conn + finally: + conn.close() + +@contextmanager +def transaction(): + with get_db_connection() as conn: + try: + yield conn + conn.commit() + except Exception: + conn.rollback() + raise + +def execute_query(query, params=None, conn=None): + if conn: + cursor = conn.cursor() + if params: + cursor.execute(query, params) + else: + cursor.execute(query) + return cursor.fetchall() + else: + with get_db_connection() as conn: + cursor = conn.cursor() + if params: + cursor.execute(query, params) + else: + cursor.execute(query) + conn.commit() + return cursor.fetchall() + +def create_tables(): + with get_db_connection() as conn: + conn.executescript(SCHEMA_SQL) + logger.info("All RAG QA Chat tables created successfully") + +# Initialize the database +create_tables() + +# +# End of Setup +############################################################ + + +############################################################ +# +# Keyword-related functions + +# Input validation +def validate_keyword(keyword): + if not isinstance(keyword, str): + raise ValueError("Keyword must be a string") + if not keyword.strip(): + raise ValueError("Keyword cannot be empty or just whitespace") + if len(keyword) > 100: + raise ValueError("Keyword is too long (max 100 characters)") + if not re.match(r'^[a-zA-Z0-9\s\-_]+$', keyword): + raise ValueError("Keyword contains invalid characters") + return keyword.strip() + +def validate_collection_name(name): + if not isinstance(name, str): + raise ValueError("Collection name must be a string") + if not name.strip(): + raise ValueError("Collection name cannot be empty or just whitespace") + if len(name) > 100: + raise ValueError("Collection name is too long (max 100 characters)") + if not re.match(r'^[a-zA-Z0-9\s\-_]+$', name): + raise ValueError("Collection name contains invalid characters") + return name.strip() + +# Core functions +def add_keyword(keyword, conn=None): + try: + validated_keyword = validate_keyword(keyword) + query = "INSERT OR IGNORE INTO rag_qa_keywords (keyword) VALUES (?)" + execute_query(query, (validated_keyword,), conn) + logger.info(f"Keyword '{validated_keyword}' added successfully") + except ValueError as e: + logger.error(f"Invalid keyword: {e}") + raise + except Exception as e: + logger.error(f"Error adding keyword '{keyword}': {e}") + raise + +def create_keyword_collection(name, parent_id=None): + try: + validated_name = validate_collection_name(name) + query = "INSERT INTO rag_qa_keyword_collections (name, parent_id) VALUES (?, ?)" + execute_query(query, (validated_name, parent_id)) + logger.info(f"Keyword collection '{validated_name}' created successfully") + except ValueError as e: + logger.error(f"Invalid collection name: {e}") + raise + except Exception as e: + logger.error(f"Error creating keyword collection '{name}': {e}") + raise + +def add_keyword_to_collection(collection_name, keyword): + try: + validated_collection_name = validate_collection_name(collection_name) + validated_keyword = validate_keyword(keyword) + + with transaction() as conn: + add_keyword(validated_keyword, conn) + + query = ''' + INSERT INTO rag_qa_collection_keywords (collection_id, keyword_id) + SELECT c.id, k.id + FROM rag_qa_keyword_collections c, rag_qa_keywords k + WHERE c.name = ? AND k.keyword = ? + ''' + execute_query(query, (validated_collection_name, validated_keyword), conn) + + logger.info(f"Keyword '{validated_keyword}' added to collection '{validated_collection_name}' successfully") + except ValueError as e: + logger.error(f"Invalid input: {e}") + raise + except Exception as e: + logger.error(f"Error adding keyword '{keyword}' to collection '{collection_name}': {e}") + raise + +def add_keywords_to_conversation(conversation_id, keywords): + if not isinstance(keywords, (list, tuple)): + raise ValueError("Keywords must be a list or tuple") + try: + with transaction() as conn: + for keyword in keywords: + validated_keyword = validate_keyword(keyword) + add_keyword(validated_keyword, conn) + + query = ''' + INSERT INTO rag_qa_conversation_keywords (conversation_id, keyword_id) + SELECT ?, id FROM rag_qa_keywords WHERE keyword = ? + ''' + execute_query(query, (conversation_id, validated_keyword), conn) + + logger.info(f"Keywords added to conversation '{conversation_id}' successfully") + except ValueError as e: + logger.error(f"Invalid keyword: {e}") + raise + except Exception as e: + logger.error(f"Error adding keywords to conversation '{conversation_id}': {e}") + raise + +def get_keywords_for_conversation(conversation_id): + try: + query = ''' + SELECT k.keyword + FROM rag_qa_keywords k + JOIN rag_qa_conversation_keywords ck ON k.id = ck.keyword_id + WHERE ck.conversation_id = ? + ''' + result = execute_query(query, (conversation_id,)) + keywords = [row[0] for row in result] + logger.info(f"Retrieved {len(keywords)} keywords for conversation '{conversation_id}'") + return keywords + except Exception as e: + logger.error(f"Error getting keywords for conversation '{conversation_id}': {e}") + raise + +def get_keywords_for_collection(collection_name): + try: + query = ''' + SELECT k.keyword + FROM rag_qa_keywords k + JOIN rag_qa_collection_keywords ck ON k.id = ck.keyword_id + JOIN rag_qa_keyword_collections c ON ck.collection_id = c.id + WHERE c.name = ? + ''' + result = execute_query(query, (collection_name,)) + keywords = [row[0] for row in result] + logger.info(f"Retrieved {len(keywords)} keywords for collection '{collection_name}'") + return keywords + except Exception as e: + logger.error(f"Error getting keywords for collection '{collection_name}': {e}") + raise + +# +# End of Keyword-related functions +################################################### + + +################################################### +# +# Notes and chat-related functions + +def save_notes(conversation_id, title, content): + """Save notes to the database.""" + try: + query = "INSERT INTO rag_qa_notes (conversation_id, title, content, timestamp) VALUES (?, ?, ?, ?)" + timestamp = datetime.now().isoformat() + with transaction() as conn: + cursor = conn.cursor() + cursor.execute(query, (conversation_id, title, content, timestamp)) + note_id = cursor.lastrowid + logger.info(f"Notes saved for conversation '{conversation_id}', note ID '{note_id}'") + return note_id + except Exception as e: + logger.error(f"Error saving notes for conversation '{conversation_id}': {e}") + raise + +def update_note(note_id, title, content): + try: + query = "UPDATE rag_qa_notes SET title = ?, content = ?, timestamp = ? WHERE id = ?" + timestamp = datetime.now().isoformat() + execute_query(query, (title, content, timestamp, note_id)) + logger.info(f"Note ID '{note_id}' updated successfully") + except Exception as e: + logger.error(f"Error updating note ID '{note_id}': {e}") + raise + +def get_notes(conversation_id): + """Retrieve notes for a given conversation.""" + try: + query = "SELECT content FROM rag_qa_notes WHERE conversation_id = ?" + result = execute_query(query, (conversation_id,)) + notes = [row[0] for row in result] + logger.info(f"Retrieved {len(notes)} notes for conversation '{conversation_id}'") + return notes + except Exception as e: + logger.error(f"Error getting notes for conversation '{conversation_id}': {e}") + raise + +def get_note_by_id(note_id): + try: + query = "SELECT id, title, content FROM rag_qa_notes WHERE id = ?" + result = execute_query(query, (note_id,)) + return result + except Exception as e: + logger.error(f"Error getting note by ID '{note_id}': {e}") + raise + +def get_notes_by_keywords(keywords, page=1, page_size=20): + try: + placeholders = ','.join(['?'] * len(keywords)) + query = f''' + SELECT n.id, n.title, n.content, n.timestamp + FROM rag_qa_notes n + JOIN rag_qa_note_keywords nk ON n.id = nk.note_id + JOIN rag_qa_keywords k ON nk.keyword_id = k.id + WHERE k.keyword IN ({placeholders}) + ORDER BY n.timestamp DESC + ''' + results, total_pages, total_count = get_paginated_results(query, tuple(keywords), page, page_size) + logger.info(f"Retrieved {len(results)} notes matching keywords: {', '.join(keywords)} (page {page} of {total_pages})") + notes = [(row[0], row[1], row[2], row[3]) for row in results] + return notes, total_pages, total_count + except Exception as e: + logger.error(f"Error getting notes by keywords: {e}") + raise + +def get_notes_by_keyword_collection(collection_name, page=1, page_size=20): + try: + query = ''' + SELECT n.id, n.title, n.content, n.timestamp + FROM rag_qa_notes n + JOIN rag_qa_note_keywords nk ON n.id = nk.note_id + JOIN rag_qa_keywords k ON nk.keyword_id = k.id + JOIN rag_qa_collection_keywords ck ON k.id = ck.keyword_id + JOIN rag_qa_keyword_collections c ON ck.collection_id = c.id + WHERE c.name = ? + ORDER BY n.timestamp DESC + ''' + results, total_pages, total_count = get_paginated_results(query, (collection_name,), page, page_size) + logger.info(f"Retrieved {len(results)} notes for collection '{collection_name}' (page {page} of {total_pages})") + notes = [(row[0], row[1], row[2], row[3]) for row in results] + return notes, total_pages, total_count + except Exception as e: + logger.error(f"Error getting notes by keyword collection '{collection_name}': {e}") + raise + +def clear_notes(conversation_id): + """Clear all notes for a given conversation.""" + try: + query = "DELETE FROM rag_qa_notes WHERE conversation_id = ?" + execute_query(query, (conversation_id,)) + logger.info(f"Cleared notes for conversation '{conversation_id}'") + except Exception as e: + logger.error(f"Error clearing notes for conversation '{conversation_id}': {e}") + raise + +def add_keywords_to_note(note_id, keywords): + """Associate keywords with a note.""" + try: + with transaction() as conn: + for keyword in keywords: + validated_keyword = validate_keyword(keyword) + add_keyword(validated_keyword, conn) + + # Retrieve the keyword ID + query = "SELECT id FROM rag_qa_keywords WHERE keyword = ?" + result = execute_query(query, (validated_keyword,), conn) + if result: + keyword_id = result[0][0] + else: + raise Exception(f"Keyword '{validated_keyword}' not found after insertion") + + # Link the note and keyword + query = "INSERT INTO rag_qa_note_keywords (note_id, keyword_id) VALUES (?, ?)" + execute_query(query, (note_id, keyword_id), conn) + + logger.info(f"Keywords added to note ID '{note_id}' successfully") + except Exception as e: + logger.error(f"Error adding keywords to note ID '{note_id}': {e}") + raise + +def get_keywords_for_note(note_id): + """Retrieve keywords associated with a given note.""" + try: + query = ''' + SELECT k.keyword + FROM rag_qa_keywords k + JOIN rag_qa_note_keywords nk ON k.id = nk.keyword_id + WHERE nk.note_id = ? + ''' + result = execute_query(query, (note_id,)) + keywords = [row[0] for row in result] + logger.info(f"Retrieved {len(keywords)} keywords for note ID '{note_id}'") + return keywords + except Exception as e: + logger.error(f"Error getting keywords for note ID '{note_id}': {e}") + raise + +def clear_keywords_from_note(note_id): + """Clear all keywords from a given note.""" + try: + query = "DELETE FROM rag_qa_note_keywords WHERE note_id = ?" + execute_query(query, (note_id,)) + logger.info(f"Cleared keywords for note ID '{note_id}'") + except Exception as e: + logger.error(f"Error clearing keywords for note ID '{note_id}': {e}") + raise + +def delete_note_by_id(note_id, conn=None): + """Delete a note and its associated keywords.""" + try: + # Delete note keywords + execute_query("DELETE FROM rag_qa_note_keywords WHERE note_id = ?", (note_id,), conn) + # Delete the note + execute_query("DELETE FROM rag_qa_notes WHERE id = ?", (note_id,), conn) + logging.info(f"Note ID '{note_id}' deleted successfully.") + except Exception as e: + logger.error(f"Error deleting note ID '{note_id}': {e}") + raise + +def delete_note(note_id): + """Delete a note by ID.""" + try: + with transaction() as conn: + delete_note_by_id(note_id, conn) + except Exception as e: + logger.error(f"Error deleting note ID '{note_id}': {e}") + raise + +# +# End of Notes related functions +################################################### + + +################################################### +# +# Chat-related functions + +def save_message(conversation_id, role, content): + try: + timestamp = datetime.now().isoformat() + query = "INSERT INTO rag_qa_chats (conversation_id, timestamp, role, content) VALUES (?, ?, ?, ?)" + execute_query(query, (conversation_id, timestamp, role, content)) + + # Update last_updated in conversation_metadata + update_query = "UPDATE conversation_metadata SET last_updated = ? WHERE conversation_id = ?" + execute_query(update_query, (timestamp, conversation_id)) + + logger.info(f"Message saved for conversation '{conversation_id}'") + except Exception as e: + logger.error(f"Error saving message for conversation '{conversation_id}': {e}") + raise + +def start_new_conversation(title="Untitled Conversation"): + try: + conversation_id = str(uuid.uuid4()) + query = "INSERT INTO conversation_metadata (conversation_id, created_at, last_updated, title) VALUES (?, ?, ?, ?)" + now = datetime.now().isoformat() + execute_query(query, (conversation_id, now, now, title)) + logger.info(f"New conversation '{conversation_id}' started with title '{title}'") + return conversation_id + except Exception as e: + logger.error(f"Error starting new conversation: {e}") + raise + +def get_all_conversations(page=1, page_size=20): + try: + query = "SELECT conversation_id, title FROM conversation_metadata ORDER BY last_updated DESC" + results, total_pages, total_count = get_paginated_results(query, page=page, page_size=page_size) + conversations = [(row[0], row[1]) for row in results] + logger.info(f"Retrieved {len(conversations)} conversations (page {page} of {total_pages})") + return conversations, total_pages, total_count + except Exception as e: + logger.error(f"Error getting conversations: {e}") + raise + +# Pagination helper function +def get_paginated_results(query, params=None, page=1, page_size=20): + try: + offset = (page - 1) * page_size + paginated_query = f"{query} LIMIT ? OFFSET ?" + if params: + paginated_params = params + (page_size, offset) + else: + paginated_params = (page_size, offset) + + result = execute_query(paginated_query, paginated_params) + + count_query = f"SELECT COUNT(*) FROM ({query}) AS total" + count_params = params if params else () + + total_count = execute_query(count_query, count_params)[0][0] + + total_pages = (total_count + page_size - 1) // page_size + + logger.info(f"Retrieved page {page} of {total_pages} (total items: {total_count})") + return result, total_pages, total_count + except Exception as e: + logger.error(f"Error retrieving paginated results: {e}") + raise + +def get_all_collections(page=1, page_size=20): + try: + query = "SELECT name FROM rag_qa_keyword_collections" + results, total_pages, total_count = get_paginated_results(query, page=page, page_size=page_size) + collections = [row[0] for row in results] + logger.info(f"Retrieved {len(collections)} keyword collections (page {page} of {total_pages})") + return collections, total_pages, total_count + except Exception as e: + logger.error(f"Error getting collections: {e}") + raise + +def search_conversations_by_keywords(keywords, page=1, page_size=20): + try: + placeholders = ','.join(['?' for _ in keywords]) + query = f''' + SELECT DISTINCT cm.conversation_id, cm.title + FROM conversation_metadata cm + JOIN rag_qa_conversation_keywords ck ON cm.conversation_id = ck.conversation_id + JOIN rag_qa_keywords k ON ck.keyword_id = k.id + WHERE k.keyword IN ({placeholders}) + ''' + results, total_pages, total_count = get_paginated_results(query, tuple(keywords), page, page_size) + logger.info( + f"Found {total_count} conversations matching keywords: {', '.join(keywords)} (page {page} of {total_pages})") + return results, total_pages, total_count + except Exception as e: + logger.error(f"Error searching conversations by keywords {keywords}: {e}") + raise + +def load_chat_history(conversation_id, page=1, page_size=50): + try: + query = "SELECT role, content FROM rag_qa_chats WHERE conversation_id = ? ORDER BY timestamp" + results, total_pages, total_count = get_paginated_results(query, (conversation_id,), page, page_size) + logger.info( + f"Loaded {len(results)} messages for conversation '{conversation_id}' (page {page} of {total_pages})") + return results, total_pages, total_count + except Exception as e: + logger.error(f"Error loading chat history for conversation '{conversation_id}': {e}") + raise + +def update_conversation_title(conversation_id, new_title): + """Update the title of a conversation.""" + try: + query = "UPDATE conversation_metadata SET title = ? WHERE conversation_id = ?" + execute_query(query, (new_title, conversation_id)) + logger.info(f"Conversation '{conversation_id}' title updated to '{new_title}'") + except Exception as e: + logger.error(f"Error updating conversation title: {e}") + raise + +def delete_conversation(conversation_id): + """Delete a conversation and its associated messages and notes.""" + try: + with transaction() as conn: + # Delete messages + execute_query("DELETE FROM rag_qa_chats WHERE conversation_id = ?", (conversation_id,), conn) + # Delete conversation metadata + execute_query("DELETE FROM conversation_metadata WHERE conversation_id = ?", (conversation_id,), conn) + # Delete conversation keywords + execute_query("DELETE FROM rag_qa_conversation_keywords WHERE conversation_id = ?", (conversation_id,), conn) + # Delete notes associated with the conversation + note_ids = execute_query("SELECT id FROM rag_qa_notes WHERE conversation_id = ?", (conversation_id,), conn) + for (note_id,) in note_ids: + delete_note_by_id(note_id, conn) + logging.info(f"Conversation '{conversation_id}' deleted successfully.") + except Exception as e: + logger.error(f"Error deleting conversation '{conversation_id}': {e}") + raise + +# +# End of Chat-related functions +################################################### + + +################################################### +# +# Functions to export DB data + +def fetch_all_conversations(): + try: + # Fetch all conversation IDs and titles + query = "SELECT conversation_id, title FROM conversation_metadata ORDER BY last_updated DESC" + results = execute_query(query) + conversations = [] + for row in results: + conversation_id, title = row + # Fetch all messages for this conversation + messages = load_all_chat_history(conversation_id) + conversations.append((conversation_id, title, messages)) + logger.info(f"Fetched all conversations: {len(conversations)} found.") + return conversations + except Exception as e: + logger.error(f"Error fetching all conversations: {e}") + raise + +def load_all_chat_history(conversation_id): + try: + query = "SELECT role, content FROM rag_qa_chats WHERE conversation_id = ? ORDER BY timestamp" + results = execute_query(query, (conversation_id,)) + messages = [(row[0], row[1]) for row in results] + return messages + except Exception as e: + logger.error(f"Error loading chat history for conversation '{conversation_id}': {e}") + raise + +def fetch_all_notes(): + try: + query = "SELECT id, title, content FROM rag_qa_notes ORDER BY timestamp DESC" + results = execute_query(query) + notes = [(row[0], row[1], row[2]) for row in results] + logger.info(f"Fetched all notes: {len(notes)} found.") + return notes + except Exception as e: + logger.error(f"Error fetching all notes: {e}") + raise + +def fetch_conversations_by_ids(conversation_ids): + try: + if not conversation_ids: + return [] + placeholders = ','.join(['?'] * len(conversation_ids)) + query = f"SELECT conversation_id, title FROM conversation_metadata WHERE conversation_id IN ({placeholders})" + results = execute_query(query, conversation_ids) + conversations = [] + for row in results: + conversation_id, title = row + # Fetch all messages for this conversation + messages = load_all_chat_history(conversation_id) + conversations.append((conversation_id, title, messages)) + logger.info(f"Fetched {len(conversations)} conversations by IDs.") + return conversations + except Exception as e: + logger.error(f"Error fetching conversations by IDs: {e}") + raise + +def fetch_notes_by_ids(note_ids): + try: + if not note_ids: + return [] + placeholders = ','.join(['?'] * len(note_ids)) + query = f"SELECT id, title, content FROM rag_qa_notes WHERE id IN ({placeholders})" + results = execute_query(query, note_ids) + notes = [(row[0], row[1], row[2]) for row in results] + logger.info(f"Fetched {len(notes)} notes by IDs.") + return notes + except Exception as e: + logger.error(f"Error fetching notes by IDs: {e}") + raise + +# +# End of Export functions +################################################### + +# +# End of RAG_QA_Chat_DB.py +#################################################################################################### diff --git a/App_Function_Libraries/DB/SQLite_DB.py b/App_Function_Libraries/DB/SQLite_DB.py new file mode 100644 index 0000000000000000000000000000000000000000..2c05cbb86041120cc0c277aebeb2e45e42eb9050 --- /dev/null +++ b/App_Function_Libraries/DB/SQLite_DB.py @@ -0,0 +1,3090 @@ +# SQLite_DB.py +######################################### +# SQLite_DB Library +# This library is used to perform any/all DB operations related to SQLite. +# +#### +import configparser +#################### +# Function List +# FIXME - UPDATE Function Arguments +# 1. get_connection(self) +# 2. execute_query(self, query: str, params: Tuple = ()) +# 3. create_tables() +# 4. add_keyword(keyword: str) +# 5. delete_keyword(keyword: str) +# 6. add_media_with_keywords(url, title, media_type, content, keywords, prompt, summary, transcription_model, author, ingestion_date) +# 7. fetch_all_keywords() +# 8. keywords_browser_interface() +# 9. display_keywords() +# 10. export_keywords_to_csv() +# 11. browse_items(search_query, search_type) +# 12. fetch_item_details(media_id: int) +# 13. add_media_version(media_id: int, prompt: str, summary: str) +# 14. search_db(search_query: str, search_fields: List[str], keywords: str, page: int = 1, results_per_page: int = 10) +# 15. search_and_display(search_query, search_fields, keywords, page) +# 16. display_details(index, results) +# 17. get_details(index, dataframe) +# 18. format_results(results) +# 19. export_to_csv(search_query: str, search_fields: List[str], keyword: str, page: int = 1, results_per_file: int = 1000) +# 20. is_valid_url(url: str) -> bool +# 21. is_valid_date(date_string: str) -> bool +# 22. add_media_to_database(url, info_dict, segments, summary, keywords, custom_prompt_input, whisper_model) +# 23. create_prompts_db() +# 24. add_prompt(name, details, system, user=None) +# 25. fetch_prompt_details(name) +# 26. list_prompts() +# 27. insert_prompt_to_db(title, description, system_prompt, user_prompt) +# 28. update_media_content(media_id: int, content: str, prompt: str, summary: str) +# 29. search_media_database(query: str) -> List[Tuple[int, str, str]] +# 30. load_media_content(media_id: int) +# 31. +# 32. +# +# +##################### +# +# Import necessary libraries +import csv +import hashlib +import html +import logging +import os +import queue +import re +import shutil +import sqlite3 +import threading +import traceback +from contextlib import contextmanager +from datetime import datetime, timedelta +from typing import List, Tuple, Dict, Any, Optional +from urllib.parse import quote + +# Local Libraries +from App_Function_Libraries.Utils.Utils import get_project_relative_path, get_database_path, \ + get_database_dir +from App_Function_Libraries.Chunk_Lib import chunk_options, chunk_text +# +# Third-Party Libraries +import gradio as gr +import pandas as pd +import yaml +# +####################################################################################################################### +# Function Definitions +# + +def ensure_database_directory(): + os.makedirs(get_database_dir(), exist_ok=True) + +ensure_database_directory() + +# Set up logging +logger = logging.getLogger(__name__) + +# FIXME - Setup properly and test/add documentation for its existence... +# Construct the path to the config file +config_path = get_project_relative_path('Config_Files/config.txt') + +# Read the config file +config = configparser.ConfigParser() +config.read(config_path) + +# Get the SQLite path from the config, or use the default if not specified +sqlite_path = config.get('Database', 'sqlite_path', fallback=get_database_path('media_summary.db')) + +# Get the backup path from the config, or use the default if not specified +backup_path = config.get('Database', 'backup_path', fallback='database_backups') +backup_path = get_project_relative_path(backup_path) + +# Set the final paths +db_path = sqlite_path +backup_dir = backup_path + +print(f"Media Database path: {db_path}") +print(f"Media Backup directory: {backup_dir}") +#create_automated_backup(db_path, backup_dir) + +# FIXME - Setup properly and test/add documentation for its existence... +#backup_file = create_automated_backup(db_path, backup_dir) +#upload_to_s3(backup_file, 'your-s3-bucket-name', f"database_backups/{os.path.basename(backup_file)}") + +# FIXME - Setup properly and test/add documentation for its existence... +#create_incremental_backup(db_path, backup_dir) + +# FIXME - Setup properly and test/add documentation for its existence... +#rotate_backups(backup_dir) + +# +# +####################################################################################################################### + + +####################################################################################################################### +# +# Backup-related functions + +def create_incremental_backup(db_path, backup_dir): + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + + # Get the page count of the database + cursor.execute("PRAGMA page_count") + page_count = cursor.fetchone()[0] + + # Create a new backup file + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + backup_file = os.path.join(backup_dir, f"incremental_backup_{timestamp}.sqlib") + + # Perform the incremental backup + conn.execute(f"VACUUM INTO '{backup_file}'") + + conn.close() + print(f"Incremental backup created: {backup_file}") + return backup_file + + +def create_automated_backup(db_path, backup_dir): + # Ensure backup directory exists + os.makedirs(backup_dir, exist_ok=True) + + # Create a timestamped backup file name + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + backup_file = os.path.join(backup_dir, f"media_db_backup_{timestamp}.db") + + # Copy the database file + shutil.copy2(db_path, backup_file) + + print(f"Backup created: {backup_file}") + return backup_file + +# FIXME - boto3 aint getting installed by default.... +# def upload_to_s3(file_path, bucket_name, s3_key): +# import boto3 +# s3 = boto3.client('s3') +# try: +# s3.upload_file(file_path, bucket_name, s3_key) +# print(f"File uploaded to S3: {s3_key}") +# except Exception as e: +# print(f"Error uploading to S3: {str(e)}") + + +def rotate_backups(backup_dir, max_backups=10): + backups = sorted( + [f for f in os.listdir(backup_dir) if f.endswith('.db')], + key=lambda x: os.path.getmtime(os.path.join(backup_dir, x)), + reverse=True + ) + + while len(backups) > max_backups: + old_backup = backups.pop() + os.remove(os.path.join(backup_dir, old_backup)) + print(f"Removed old backup: {old_backup}") + +# +# +####################################################################################################################### + + +####################################################################################################################### +# +# DB-Integrity Check Functions + +def check_database_integrity(db_path): + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + + cursor.execute("PRAGMA integrity_check") + result = cursor.fetchone() + + conn.close() + + if result[0] == "ok": + print("Database integrity check passed.") + return True + else: + print("Database integrity check failed:", result[0]) + return False + +#check_database_integrity(db_path) + +# +# End of DB-Integrity Check functions +####################################################################################################################### + + +####################################################################################################################### +# +# DB Setup Functions + +class DatabaseError(Exception): + pass + +class InputError(Exception): + pass + + +class Database: + def __init__(self, db_name='media_summary.db'): + self.db_path = get_database_path(db_name) + self.timeout = 10.0 + self._local = threading.local() + + @contextmanager + def get_connection(self): + if not hasattr(self._local, 'connection') or self._local.connection is None: + self._local.connection = sqlite3.connect(self.db_path, timeout=self.timeout) + self._local.connection.isolation_level = None # This enables autocommit mode + yield self._local.connection + + def close_connection(self): + if hasattr(self._local, 'connection') and self._local.connection: + self._local.connection.close() + self._local.connection = None + + @contextmanager + def transaction(self): + with self.get_connection() as conn: + try: + conn.execute("BEGIN") + yield conn + conn.execute("COMMIT") + except Exception: + conn.execute("ROLLBACK") + raise + + def execute_query(self, query: str, params: Tuple = ()) -> Any: + with self.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(query, params) + if query.strip().upper().startswith("SELECT"): + return cursor.fetchall() + else: + return cursor.rowcount + + def execute_many(self, query: str, params_list: List[Tuple]) -> None: + with self.get_connection() as conn: + cursor = conn.cursor() + cursor.executemany(query, params_list) + + def table_exists(self, table_name: str) -> bool: + query = 'SELECT name FROM sqlite_master WHERE type="table" AND name=?' + result = self.execute_query(query, (table_name,)) + return bool(result) + +db = Database() + +# Usage example: +if db.table_exists('DocumentVersions'): + logging.debug("DocumentVersions table exists") +else: + logging.debug("DocumentVersions table does not exist") + + +# Function to create tables with the new media schema +def create_tables(db) -> None: + table_queries = [ + # CREATE TABLE statements + ''' + CREATE TABLE IF NOT EXISTS Media ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + url TEXT, + title TEXT NOT NULL, + type TEXT NOT NULL, + content TEXT, + author TEXT, + ingestion_date TEXT, + prompt TEXT, + summary TEXT, + transcription_model TEXT, + is_trash BOOLEAN DEFAULT 0, + trash_date DATETIME, + vector_embedding BLOB, + chunking_status TEXT DEFAULT 'pending', + vector_processing INTEGER DEFAULT 0 + ) + ''', + ''' + CREATE TABLE IF NOT EXISTS Keywords ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + keyword TEXT NOT NULL UNIQUE + ) + ''', + ''' + CREATE TABLE IF NOT EXISTS MediaKeywords ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + media_id INTEGER NOT NULL, + keyword_id INTEGER NOT NULL, + FOREIGN KEY (media_id) REFERENCES Media(id), + FOREIGN KEY (keyword_id) REFERENCES Keywords(id) + ) + ''', + ''' + CREATE TABLE IF NOT EXISTS MediaVersion ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + media_id INTEGER NOT NULL, + version INTEGER NOT NULL, + prompt TEXT, + summary TEXT, + created_at TEXT NOT NULL, + FOREIGN KEY (media_id) REFERENCES Media(id) + ) + ''', + ''' + CREATE TABLE IF NOT EXISTS MediaModifications ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + media_id INTEGER NOT NULL, + prompt TEXT, + summary TEXT, + modification_date TEXT, + FOREIGN KEY (media_id) REFERENCES Media(id) + ) + ''', + ''' + CREATE TABLE IF NOT EXISTS ChatConversations ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + media_id INTEGER, + media_name TEXT, + conversation_name TEXT, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (media_id) REFERENCES Media(id) + ) + ''', + ''' + CREATE TABLE IF NOT EXISTS ChatMessages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + conversation_id INTEGER, + sender TEXT, + message TEXT, + timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (conversation_id) REFERENCES ChatConversations(id) + ) + ''', + ''' + CREATE TABLE IF NOT EXISTS Transcripts ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + media_id INTEGER, + whisper_model TEXT, + transcription TEXT, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (media_id) REFERENCES Media(id) + ) + ''', + ''' + CREATE TABLE IF NOT EXISTS MediaChunks ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + media_id INTEGER, + chunk_text TEXT, + start_index INTEGER, + end_index INTEGER, + chunk_id TEXT, + FOREIGN KEY (media_id) REFERENCES Media(id) + )''', + ''' + CREATE TABLE IF NOT EXISTS UnvectorizedMediaChunks ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + media_id INTEGER NOT NULL, + chunk_text TEXT NOT NULL, + chunk_index INTEGER NOT NULL, + start_char INTEGER NOT NULL, + end_char INTEGER NOT NULL, + chunk_type TEXT, + creation_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + last_modified TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + is_processed BOOLEAN DEFAULT FALSE, + metadata TEXT, + FOREIGN KEY (media_id) REFERENCES Media(id) + ) + ''', + ''' + CREATE TABLE IF NOT EXISTS DocumentVersions ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + media_id INTEGER NOT NULL, + version_number INTEGER NOT NULL, + content TEXT, + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (media_id) REFERENCES Media(id) + ) + ''', + ] + + index_queries = [ + # CREATE INDEX statements + 'CREATE INDEX IF NOT EXISTS idx_media_title ON Media(title)', + 'CREATE INDEX IF NOT EXISTS idx_media_type ON Media(type)', + 'CREATE INDEX IF NOT EXISTS idx_media_author ON Media(author)', + 'CREATE INDEX IF NOT EXISTS idx_media_ingestion_date ON Media(ingestion_date)', + 'CREATE INDEX IF NOT EXISTS idx_keywords_keyword ON Keywords(keyword)', + 'CREATE INDEX IF NOT EXISTS idx_mediakeywords_media_id ON MediaKeywords(media_id)', + 'CREATE INDEX IF NOT EXISTS idx_mediakeywords_keyword_id ON MediaKeywords(keyword_id)', + 'CREATE INDEX IF NOT EXISTS idx_media_version_media_id ON MediaVersion(media_id)', + 'CREATE INDEX IF NOT EXISTS idx_mediamodifications_media_id ON MediaModifications(media_id)', + 'CREATE INDEX IF NOT EXISTS idx_chatconversations_media_id ON ChatConversations(media_id)', + 'CREATE INDEX IF NOT EXISTS idx_chatmessages_conversation_id ON ChatMessages(conversation_id)', + 'CREATE INDEX IF NOT EXISTS idx_media_is_trash ON Media(is_trash)', + 'CREATE INDEX IF NOT EXISTS idx_mediachunks_media_id ON MediaChunks(media_id)', + 'CREATE INDEX IF NOT EXISTS idx_unvectorized_media_chunks_media_id ON UnvectorizedMediaChunks(media_id)', + 'CREATE INDEX IF NOT EXISTS idx_unvectorized_media_chunks_is_processed ON UnvectorizedMediaChunks(is_processed)', + 'CREATE INDEX IF NOT EXISTS idx_unvectorized_media_chunks_chunk_type ON UnvectorizedMediaChunks(chunk_type)', + 'CREATE UNIQUE INDEX IF NOT EXISTS idx_unique_media_url ON Media(url)', + 'CREATE UNIQUE INDEX IF NOT EXISTS idx_unique_media_keyword ON MediaKeywords(media_id, keyword_id)', + 'CREATE INDEX IF NOT EXISTS idx_document_versions_media_id ON DocumentVersions(media_id)', + 'CREATE INDEX IF NOT EXISTS idx_document_versions_version_number ON DocumentVersions(version_number)', + ] + + virtual_table_queries = [ + # CREATE VIRTUAL TABLE statements + 'CREATE VIRTUAL TABLE IF NOT EXISTS media_fts USING fts5(title, content)', + 'CREATE VIRTUAL TABLE IF NOT EXISTS keyword_fts USING fts5(keyword)' + ] + + all_queries = table_queries + index_queries + virtual_table_queries + + for query in all_queries: + try: + db.execute_query(query) + except Exception as e: + logging.error(f"Error executing query: {query}") + logging.error(f"Error details: {str(e)}") + raise + + logging.info("All tables, indexes, and virtual tables created successfully.") + +create_tables(db) + +# +# End of DB Setup Functions +####################################################################################################################### + + +####################################################################################################################### +# +# Media-related Functions + +def check_media_exists(title: str, url: str) -> Optional[int]: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + query = 'SELECT id FROM Media WHERE title = ? OR url = ?' + cursor.execute(query, (title, url)) + result = cursor.fetchone() + logging.debug(f"check_media_exists query: {query}") + logging.debug(f"check_media_exists params: title={title}, url={url}") + logging.debug(f"check_media_exists result: {result}") + return result[0] if result else None + except Exception as e: + logging.error(f"Error checking if media exists: {str(e)}") + logging.error(f"Exception details: {traceback.format_exc()}") + return None + + +def check_media_and_whisper_model(title=None, url=None, current_whisper_model=None): + """ + Check if media exists in the database and compare the whisper model used. + + :param title: Title of the media (optional) + :param url: URL of the media (optional) + :param current_whisper_model: The whisper model currently selected for use + :return: Tuple (bool, str) - (should_download, reason) + """ + if not title and not url: + return True, "No title or URL provided" + + with db.get_connection() as conn: + cursor = conn.cursor() + + # First, find the media_id + query = "SELECT id FROM Media WHERE " + params = [] + + if title: + query += "title = ?" + params.append(title) + + if url: + if params: + query += " OR " + query += "url = ?" + params.append(url) + + cursor.execute(query, tuple(params)) + result = cursor.fetchone() + + if not result: + return True, "Media not found in database" + + media_id = result[0] + + # Now, get the latest transcript for this media + cursor.execute(""" + SELECT transcription + FROM Transcripts + WHERE media_id = ? + ORDER BY created_at DESC + LIMIT 1 + """, (media_id,)) + + transcript_result = cursor.fetchone() + + if not transcript_result: + return True, f"No transcript found for media (ID: {media_id})" + + transcription = transcript_result[0] + + # Extract the whisper model from the transcription + match = re.search(r"This text was transcribed using whisper model: (.+)$", transcription, re.MULTILINE) + if not match: + return True, f"Whisper model information not found in transcript (Media ID: {media_id})" + + db_whisper_model = match.group(1).strip() + + if not current_whisper_model: + return False, f"Media found in database (ID: {media_id})" + + if db_whisper_model != current_whisper_model: + return True, f"Different whisper model (DB: {db_whisper_model}, Current: {current_whisper_model})" + + return False, f"Media found with same whisper model (ID: {media_id})" + + +def add_media_chunk(media_id: int, chunk_text: str, start_index: int, end_index: int, chunk_id: str): + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute( + "INSERT INTO MediaChunks (media_id, chunk_text, start_index, end_index, chunk_id) VALUES (?, ?, ?, ?, ?)", + (media_id, chunk_text, start_index, end_index, chunk_id) + ) + conn.commit() + +def sqlite_update_fts_for_media(db, media_id: int): + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute("INSERT OR REPLACE INTO media_fts (rowid, title, content) SELECT id, title, content FROM Media WHERE id = ?", (media_id,)) + conn.commit() + + +def get_unprocessed_media(db): + query = """ + SELECT id, content, type, COALESCE(title, '') as file_name + FROM Media + WHERE vector_processing = 0 + ORDER BY id + """ + return db.execute_query(query) + +def get_next_media_id(): + try: + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + cursor.execute("SELECT MAX(media_id) FROM media") + max_id = cursor.fetchone()[0] + return (max_id or 0) + 1 + finally: + conn.close() + + +def mark_media_as_processed(database, media_id): + try: + query = "UPDATE Media SET vector_processing = 1 WHERE id = ?" + database.execute_query(query, (media_id,)) + logger.info(f"Marked media_id {media_id} as processed") + except Exception as e: + logger.error(f"Error marking media_id {media_id} as processed: {str(e)}") + raise + +# +# End of Vector-chunk-related Functions +####################################################################################################################### + + +####################################################################################################################### +# Keyword-related Functions +# + +# Function to add media with keywords +def add_media_with_keywords(url, title, media_type, content, keywords, prompt, summary, transcription_model, author, + ingestion_date): + logging.debug(f"Entering add_media_with_keywords: URL={url}, Title={title}") + # Set default values for missing fields + if url is None: + url = 'localhost' + elif url is not None: + url = url + title = title or 'Untitled' + media_type = media_type or 'Unknown' + content = content or 'No content available' + keywords = keywords or 'default' + prompt = prompt or 'No prompt available' + summary = summary or 'No summary available' + transcription_model = transcription_model or 'Unknown' + author = author or 'Unknown' + ingestion_date = ingestion_date or datetime.now().strftime('%Y-%m-%d') + + if media_type not in ['article', 'audio', 'document', 'mediawiki_article', 'mediawiki_dump', 'obsidian_note', 'podcast', 'text', 'video', 'unknown']: + raise InputError("Invalid media type. Allowed types: article, audio file, document, obsidian_note podcast, text, video, unknown.") + + if ingestion_date and not is_valid_date(ingestion_date): + raise InputError("Invalid ingestion date format. Use YYYY-MM-DD.") + + # Handle keywords as either string or list + if isinstance(keywords, str): + keyword_list = [keyword.strip().lower() for keyword in keywords.split(',')] + elif isinstance(keywords, list): + keyword_list = [keyword.strip().lower() for keyword in keywords] + else: + keyword_list = ['default'] + + logging.info(f"Adding/updating media: URL={url}, Title={title}, Type={media_type}") + logging.debug(f"Content (first 500 chars): {content[:500]}...") + logging.debug(f"Keywords: {keyword_list}") + logging.info(f"Prompt: {prompt}") + logging.info(f"Summary: {summary}") + logging.info(f"Author: {author}") + logging.info(f"Ingestion Date: {ingestion_date}") + logging.info(f"Transcription Model: {transcription_model}") + + try: + with db.get_connection() as conn: + cursor = conn.cursor() + + # Check if media already exists using both title and URL + existing_media_id = check_media_exists(title, url) + logging.debug(f"Existing media ID for {url}: {existing_media_id}") + + if existing_media_id: + media_id = existing_media_id + logging.debug(f"Updating existing media with ID: {media_id}") + cursor.execute(''' + UPDATE Media + SET content = ?, transcription_model = ?, type = ?, author = ?, ingestion_date = ? + WHERE id = ? + ''', (content, transcription_model, media_type, author, ingestion_date, media_id)) + else: + logging.debug("Inserting new media") + cursor.execute(''' + INSERT INTO Media (url, title, type, content, author, ingestion_date, transcription_model) + VALUES (?, ?, ?, ?, ?, ?, ?) + ''', (url, title, media_type, content, author, ingestion_date, transcription_model)) + media_id = cursor.lastrowid + logging.debug(f"New media inserted with ID: {media_id}") + + cursor.execute(''' + INSERT INTO MediaModifications (media_id, prompt, summary, modification_date) + VALUES (?, ?, ?, ?) + ''', (media_id, prompt, summary, ingestion_date)) + + # Batch insert keywords + keyword_params = [(keyword.strip().lower(),) for keyword in keyword_list] + cursor.executemany('INSERT OR IGNORE INTO Keywords (keyword) VALUES (?)', keyword_params) + + # Get keyword IDs + placeholder = ','.join(['?'] * len(keyword_list)) + cursor.execute(f'SELECT id, keyword FROM Keywords WHERE keyword IN ({placeholder})', keyword_list) + keyword_ids = cursor.fetchall() + + # Batch insert media-keyword associations + media_keyword_params = [(media_id, keyword_id) for keyword_id, _ in keyword_ids] + cursor.executemany('INSERT OR IGNORE INTO MediaKeywords (media_id, keyword_id) VALUES (?, ?)', media_keyword_params) + + # Update full-text search index + cursor.execute('INSERT OR REPLACE INTO media_fts (rowid, title, content) VALUES (?, ?, ?)', + (media_id, title, content)) + + # Add media version + add_media_version(conn, media_id, prompt, summary) + + conn.commit() + logging.info(f"Media '{title}' successfully added/updated with ID: {media_id}") + + return media_id, f"Media '{title}' added/updated successfully with keywords: {', '.join(keyword_list)}" + + except sqlite3.Error as e: + logging.error(f"SQL Error in add_media_with_keywords: {e}") + raise DatabaseError(f"Error adding media with keywords: {e}") + except Exception as e: + logging.error(f"Unexpected Error in add_media_with_keywords: {e}") + raise DatabaseError(f"Unexpected error: {e}") + + +def ingest_article_to_db(url, title, author, content, keywords, summary, ingestion_date, custom_prompt): + try: + # Check if content is not empty or whitespace + if not content.strip(): + raise ValueError("Content is empty.") + + keyword_list = keywords.split(',') if keywords else ["default"] + keyword_str = ', '.join(keyword_list) + + # Set default values for missing fields + url = url or 'Unknown' + title = title or 'Unknown' + author = author or 'Unknown' + keywords = keywords or 'default' + summary = summary or 'No summary available' + ingestion_date = ingestion_date or datetime.now().strftime('%Y-%m-%d') + + # Log the values of all fields before calling add_media_with_keywords + logging.debug(f"URL: {url}") + logging.debug(f"Title: {title}") + logging.debug(f"Author: {author}") + logging.debug(f"Content: {content[:50]}... (length: {len(content)})") # Log first 50 characters of content + logging.debug(f"Keywords: {keywords}") + logging.debug(f"Summary: {summary}") + logging.debug(f"Ingestion Date: {ingestion_date}") + logging.debug(f"Custom Prompt: {custom_prompt}") + + # Check if any required field is empty and log the specific missing field + if not url: + logging.error("URL is missing.") + raise ValueError("URL is missing.") + if not title: + logging.error("Title is missing.") + raise ValueError("Title is missing.") + if not content: + logging.error("Content is missing.") + raise ValueError("Content is missing.") + if not keywords: + logging.error("Keywords are missing.") + raise ValueError("Keywords are missing.") + if not summary: + logging.error("Summary is missing.") + raise ValueError("Summary is missing.") + if not ingestion_date: + logging.error("Ingestion date is missing.") + raise ValueError("Ingestion date is missing.") + if not custom_prompt: + logging.error("Custom prompt is missing.") + raise ValueError("Custom prompt is missing.") + + # Add media with keywords to the database + result = add_media_with_keywords( + url=url, + title=title, + media_type='article', + content=content, + keywords=keyword_str or "article_default", + prompt=custom_prompt or None, + summary=summary or "No summary generated", + transcription_model=None, # or some default value if applicable + author=author or 'Unknown', + ingestion_date=ingestion_date + ) + return result + except Exception as e: + logging.error(f"Failed to ingest article to the database: {e}") + return str(e) + + +# Function to add a keyword +def add_keyword(keyword: str) -> int: + if not keyword.strip(): + raise DatabaseError("Keyword cannot be empty") + + keyword = keyword.strip().lower() + with db.get_connection() as conn: + cursor = conn.cursor() + try: + # Insert into Keywords table + cursor.execute('INSERT OR IGNORE INTO Keywords (keyword) VALUES (?)', (keyword,)) + + # Get the keyword_id (whether it was just inserted or already existed) + cursor.execute('SELECT id FROM Keywords WHERE keyword = ?', (keyword,)) + keyword_id = cursor.fetchone()[0] + + # Check if the keyword exists in keyword_fts + cursor.execute('SELECT rowid FROM keyword_fts WHERE rowid = ?', (keyword_id,)) + if not cursor.fetchone(): + # If it doesn't exist in keyword_fts, insert it + cursor.execute('INSERT OR IGNORE INTO keyword_fts (rowid, keyword) VALUES (?, ?)', (keyword_id, keyword)) + + logging.info(f"Keyword '{keyword}' added or updated with ID: {keyword_id}") + conn.commit() + return keyword_id + except sqlite3.IntegrityError as e: + logging.error(f"Integrity error adding keyword: {e}") + raise DatabaseError(f"Integrity error adding keyword: {e}") + except sqlite3.Error as e: + logging.error(f"Error adding keyword: {e}") + raise DatabaseError(f"Error adding keyword: {e}") + + + +# Function to delete a keyword +def delete_keyword(keyword: str) -> str: + keyword = keyword.strip().lower() + with db.get_connection() as conn: + cursor = conn.cursor() + try: + cursor.execute('SELECT id FROM Keywords WHERE keyword = ?', (keyword,)) + keyword_id = cursor.fetchone() + if keyword_id: + cursor.execute('DELETE FROM Keywords WHERE keyword = ?', (keyword,)) + cursor.execute('DELETE FROM keyword_fts WHERE rowid = ?', (keyword_id[0],)) + conn.commit() + return f"Keyword '{keyword}' deleted successfully." + else: + return f"Keyword '{keyword}' not found." + except sqlite3.Error as e: + raise DatabaseError(f"Error deleting keyword: {e}") + + +def fetch_all_keywords() -> List[str]: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute('SELECT keyword FROM Keywords') + keywords = [row[0] for row in cursor.fetchall()] + return keywords + except sqlite3.Error as e: + raise DatabaseError(f"Error fetching keywords: {e}") + +def keywords_browser_interface(): + keywords = fetch_all_keywords() + return gr.Markdown("\n".join(f"- {keyword}" for keyword in keywords)) + +def display_keywords(): + try: + keywords = fetch_all_keywords() + return "\n".join(keywords) if keywords else "No keywords found." + except DatabaseError as e: + return str(e) + + +def export_keywords_to_csv(): + try: + keywords = fetch_all_keywords() + if not keywords: + return None, "No keywords found in the database." + + filename = "keywords.csv" + with open(filename, 'w', newline='', encoding='utf-8') as file: + writer = csv.writer(file) + writer.writerow(["Keyword"]) + for keyword in keywords: + writer.writerow([keyword]) + + return filename, f"Keywords exported to {filename}" + except Exception as e: + logger.error(f"Error exporting keywords to CSV: {e}") + return None, f"Error exporting keywords: {e}" + +def fetch_keywords_for_media(media_id): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT k.keyword + FROM Keywords k + JOIN MediaKeywords mk ON k.id = mk.keyword_id + WHERE mk.media_id = ? + ''', (media_id,)) + keywords = [row[0] for row in cursor.fetchall()] + return keywords + except sqlite3.Error as e: + logging.error(f"Error fetching keywords: {e}") + return [] + +def update_keywords_for_media(media_id, keyword_list): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + + # Remove old keywords + cursor.execute('DELETE FROM MediaKeywords WHERE media_id = ?', (media_id,)) + + # Add new keywords + for keyword in keyword_list: + cursor.execute('INSERT OR IGNORE INTO Keywords (keyword) VALUES (?)', (keyword,)) + cursor.execute('SELECT id FROM Keywords WHERE keyword = ?', (keyword,)) + keyword_id = cursor.fetchone()[0] + cursor.execute('INSERT INTO MediaKeywords (media_id, keyword_id) VALUES (?, ?)', (media_id, keyword_id)) + + conn.commit() + return "Keywords updated successfully." + except sqlite3.Error as e: + logging.error(f"Error updating keywords: {e}") + return "Error updating keywords." + +# +# End of Keyword-related functions +####################################################################################################################### + + +####################################################################################################################### +# +# Media-related Functions + + + +# Function to fetch items based on search query and type +def browse_items(search_query, search_type): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + if search_type == 'Title': + cursor.execute("SELECT id, title, url FROM Media WHERE title LIKE ?", (f'%{search_query}%',)) + elif search_type == 'URL': + cursor.execute("SELECT id, title, url FROM Media WHERE url LIKE ?", (f'%{search_query}%',)) + elif search_type == 'Keyword': + return fetch_items_by_keyword(search_query) + elif search_type == 'Content': + cursor.execute("SELECT id, title, url FROM Media WHERE content LIKE ?", (f'%{search_query}%',)) + else: + raise ValueError(f"Invalid search type: {search_type}") + + results = cursor.fetchall() + return results + except sqlite3.Error as e: + logger.error(f"Error fetching items by {search_type}: {e}") + raise DatabaseError(f"Error fetching items by {search_type}: {e}") + + +# Function to fetch item details + +def fetch_item_details(media_id: int): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + # Fetch the latest prompt and summary from MediaModifications + cursor.execute(""" + SELECT prompt, summary + FROM MediaModifications + WHERE media_id = ? + ORDER BY modification_date DESC + LIMIT 1 + """, (media_id,)) + prompt_summary_result = cursor.fetchone() + + # Fetch the latest transcription + cursor.execute("SELECT content FROM Media WHERE id = ?", (media_id,)) + content_result = cursor.fetchone() + + prompt = prompt_summary_result[0] if prompt_summary_result else "No prompt available." + summary = prompt_summary_result[1] if prompt_summary_result else "No summary available." + content = content_result[0] if content_result else "No content available." + + return prompt, summary, content + except sqlite3.Error as e: + logging.error(f"Error fetching item details: {e}") + return "Error fetching prompt.", "Error fetching summary.", "Error fetching media." + +# +# End of Media-related Functions +####################################################################################################################### + + +####################################################################################################################### +# +# Media-related Functions + + +# Function to add a version of a prompt and summary +def add_media_version(conn, media_id: int, prompt: str, summary: str) -> None: + try: + cursor = conn.cursor() + + # Get the current version number + cursor.execute('SELECT MAX(version) FROM MediaVersion WHERE media_id = ?', (media_id,)) + current_version = cursor.fetchone()[0] or 0 + + # Insert the new version + cursor.execute(''' + INSERT INTO MediaVersion (media_id, version, prompt, summary, created_at) + VALUES (?, ?, ?, ?, ?) + ''', (media_id, current_version + 1, prompt, summary, datetime.now().strftime('%Y-%m-%d %H:%M:%S'))) + except DatabaseError as e: + logging.error(f"Error adding media version: {e}") + raise + + +# Function to search the database with advanced options, including keyword search and full-text search +def sqlite_search_db(search_query: str, search_fields: List[str], keywords: str, page: int = 1, results_per_page: int = 10, connection=None): + if page < 1: + raise ValueError("Page number must be 1 or greater.") + + # Prepare keywords by splitting and trimming + keywords = [keyword.strip().lower() for keyword in keywords.split(',') if keyword.strip()] + + def execute_query(conn): + cursor = conn.cursor() + offset = (page - 1) * results_per_page + + # Prepare the search conditions for general fields + search_conditions = [] + params = [] + + for field in search_fields: + if search_query: # Ensure there's a search query before adding this condition + search_conditions.append(f"Media.{field} LIKE ?") + params.append(f'%{search_query}%') + + # Prepare the conditions for keywords filtering + keyword_conditions = [] + for keyword in keywords: + keyword_conditions.append( + f"EXISTS (SELECT 1 FROM MediaKeywords mk JOIN Keywords k ON mk.keyword_id = k.id WHERE mk.media_id = Media.id AND k.keyword LIKE ?)") + params.append(f'%{keyword}%') + + # Combine all conditions + where_clause = " AND ".join( + search_conditions + keyword_conditions) if search_conditions or keyword_conditions else "1=1" + + # Complete the query + query = f''' + SELECT DISTINCT Media.id, Media.url, Media.title, Media.type, Media.content, Media.author, Media.ingestion_date, + MediaModifications.prompt, MediaModifications.summary + FROM Media + LEFT JOIN MediaModifications ON Media.id = MediaModifications.media_id + WHERE {where_clause} + ORDER BY Media.ingestion_date DESC + LIMIT ? OFFSET ? + ''' + params.extend([results_per_page, offset]) + + cursor.execute(query, params) + return cursor.fetchall() + + if connection: + return execute_query(connection) + else: + with db.get_connection() as conn: + return execute_query(conn) + + +# Gradio function to handle user input and display results with pagination, with better feedback +def search_and_display(search_query, search_fields, keywords, page): + results = sqlite_search_db(search_query, search_fields, keywords, page) + + if isinstance(results, pd.DataFrame): + # Convert DataFrame to a list of tuples or lists + processed_results = results.values.tolist() # This converts DataFrame rows to lists + elif isinstance(results, list): + # Ensure that each element in the list is itself a list or tuple (not a dictionary) + processed_results = [list(item.values()) if isinstance(item, dict) else item for item in results] + else: + raise TypeError("Unsupported data type for results") + + return processed_results + + +def display_details(index, results): + if index is None or results is None: + return "Please select a result to view details." + + try: + # Ensure the index is an integer and access the row properly + index = int(index) + if isinstance(results, pd.DataFrame): + if index >= len(results): + return "Index out of range. Please select a valid index." + selected_row = results.iloc[index] + else: + # If results is not a DataFrame, but a list (assuming list of dicts) + selected_row = results[index] + except ValueError: + return "Index must be an integer." + except IndexError: + return "Index out of range. Please select a valid index." + + # Build HTML output safely + details_html = f""" +

{selected_row.get('Title', 'No Title')}

+

URL: {selected_row.get('URL', 'No URL')}

+

Type: {selected_row.get('Type', 'No Type')}

+

Author: {selected_row.get('Author', 'No Author')}

+

Ingestion Date: {selected_row.get('Ingestion Date', 'No Date')}

+

Prompt: {selected_row.get('Prompt', 'No Prompt')}

+

Summary: {selected_row.get('Summary', 'No Summary')}

+

Content: {selected_row.get('Content', 'No Content')}

+ """ + return details_html + + +def get_details(index, dataframe): + if index is None or dataframe is None or index >= len(dataframe): + return "Please select a result to view details." + row = dataframe.iloc[index] + details = f""" +

{row['Title']}

+

URL: {row['URL']}

+

Type: {row['Type']}

+

Author: {row['Author']}

+

Ingestion Date: {row['Ingestion Date']}

+

Prompt: {row['Prompt']}

+

Summary: {row['Summary']}

+

Content:

+ 
{row['Content']}
+ """ + return details + + +def format_results(results): + if not results: + return pd.DataFrame(columns=['URL', 'Title', 'Type', 'Content', 'Author', 'Ingestion Date', 'Prompt', 'Summary']) + + df = pd.DataFrame(results, columns=['URL', 'Title', 'Type', 'Content', 'Author', 'Ingestion Date', 'Prompt', 'Summary']) + logging.debug(f"Formatted DataFrame: {df}") + + return df + + +# Function to export search results to CSV or markdown with pagination +def export_to_file(search_query: str, search_fields: List[str], keyword: str, page: int = 1, results_per_file: int = 1000, export_format: str = 'csv'): + try: + results = sqlite_search_db(search_query, search_fields, keyword, page, results_per_file) + if not results: + return "No results found to export." + + # Create an 'exports' directory if it doesn't exist + if not os.path.exists('exports'): + os.makedirs('exports') + + if export_format == 'csv': + filename = f'exports/search_results_page_{page}.csv' + with open(filename, 'w', newline='', encoding='utf-8') as file: + writer = csv.writer(file) + writer.writerow(['URL', 'Title', 'Type', 'Content', 'Author', 'Ingestion Date', 'Prompt', 'Summary']) + for row in results: + writer.writerow(row) + elif export_format == 'markdown': + filename = f'exports/search_results_page_{page}.md' + with open(filename, 'w', encoding='utf-8') as file: + for item in results: + markdown_content = convert_to_markdown({ + 'title': item[1], + 'url': item[0], + 'type': item[2], + 'content': item[3], + 'author': item[4], + 'ingestion_date': item[5], + 'summary': item[7], + 'keywords': item[8].split(',') if item[8] else [] + }) + file.write(markdown_content) + file.write("\n---\n\n") # Separator between items + else: + return f"Unsupported export format: {export_format}" + + return f"Results exported to {filename}" + except (DatabaseError, InputError) as e: + return str(e) + + +# Helper function to validate date format +def is_valid_date(date_string: str) -> bool: + try: + datetime.strptime(date_string, '%Y-%m-%d') + return True + except ValueError: + return False + + +def add_media_to_database(url, info_dict, segments, summary, keywords, custom_prompt_input, whisper_model, media_type='video', overwrite=False, db=None): + if db is None: + db = Database() + try: + with db.get_connection() as conn: + cursor = conn.cursor() + + # Generate URL if not provided + if not url: + title = info_dict.get('title', 'Untitled') + url_hash = hashlib.md5(f"{title}{media_type}".encode()).hexdigest() + url = f"https://No-URL-Submitted.com/{media_type}/{quote(title)}-{url_hash}" + + logging.debug(f"Checking for existing media with URL: {url}") + + # Extract content from segments + if isinstance(segments, list): + content = ' '.join([segment.get('Text', '') for segment in segments if 'Text' in segment]) + elif isinstance(segments, dict): + content = segments.get('text', '') or segments.get('content', '') + else: + content = str(segments) + + # Process keywords + if isinstance(keywords, str): + keyword_list = [keyword.strip().lower() for keyword in keywords.split(',')] + elif isinstance(keywords, (list, tuple)): + keyword_list = [keyword.strip().lower() for keyword in keywords] + else: + keyword_list = ['default'] + + # Check if media already exists + cursor.execute('SELECT id FROM Media WHERE url = ?', (url,)) + existing_media = cursor.fetchone() + + logging.debug(f"Existing media: {existing_media}") + logging.debug(f"Overwrite flag: {overwrite}") + + if existing_media: + media_id = existing_media[0] + logging.debug(f"Existing media_id: {media_id}") + if overwrite: + logging.debug("Updating existing media") + cursor.execute(''' + UPDATE Media + SET content = ?, transcription_model = ?, title = ?, type = ?, author = ?, ingestion_date = ?, chunking_status = ? + WHERE id = ? + ''', (content, whisper_model, info_dict.get('title', 'Untitled'), media_type, + info_dict.get('uploader', 'Unknown'), datetime.now().strftime('%Y-%m-%d'), 'pending', media_id)) + action = "updated" + else: + logging.debug("Media exists but not updating (overwrite=False)") + action = "already exists (not updated)" + else: + cursor.execute(''' + INSERT INTO Media (url, title, type, content, author, ingestion_date, transcription_model, chunking_status) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ''', (url, info_dict.get('title', 'Untitled'), media_type, content, + info_dict.get('uploader', 'Unknown'), datetime.now().strftime('%Y-%m-%d'), whisper_model, 'pending')) + media_id = cursor.lastrowid + action = "added" + logging.debug(f"New media_id: {media_id}") + + logging.debug(f"Before MediaModifications insert, media_id: {media_id}") + + # Only proceed with modifications if the media was added or updated + if action in ["updated", "added"]: + cursor.execute(''' + INSERT INTO MediaModifications (media_id, prompt, summary, modification_date) + VALUES (?, ?, ?, ?) + ''', (media_id, custom_prompt_input, summary, datetime.now().strftime('%Y-%m-%d'))) + + # Process keywords + for keyword in keyword_list: + cursor.execute('INSERT OR IGNORE INTO Keywords (keyword) VALUES (?)', (keyword,)) + cursor.execute('SELECT id FROM Keywords WHERE keyword = ?', (keyword,)) + keyword_id = cursor.fetchone()[0] + cursor.execute('INSERT OR IGNORE INTO MediaKeywords (media_id, keyword_id) VALUES (?, ?)', + (media_id, keyword_id)) + + # Update full-text search index + cursor.execute('INSERT OR REPLACE INTO media_fts (rowid, title, content) VALUES (?, ?, ?)', + (media_id, info_dict.get('title', 'Untitled'), content)) + + # Add media version + cursor.execute('SELECT MAX(version) FROM MediaVersion WHERE media_id = ?', (media_id,)) + current_version = cursor.fetchone()[0] or 0 + cursor.execute(''' + INSERT INTO MediaVersion (media_id, version, prompt, summary, created_at) + VALUES (?, ?, ?, ?, ?) + ''', (media_id, current_version + 1, custom_prompt_input, summary, datetime.now().strftime('%Y-%m-%d %H:%M:%S'))) + + conn.commit() + + # Schedule chunking + schedule_chunking(media_id, content, info_dict.get('title', 'Untitled')) + + action = "updated" if existing_media and overwrite else "added" + return f"Media '{info_dict.get('title', 'Untitled')}' {action} with URL: {url}" + \ + (f" and keywords: {', '.join(keyword_list)}. Chunking scheduled." if action in ["updated", "added"] else "") + + except DatabaseError as e: + logging.error(f"Database error: {e}") + raise + except Exception as e: + logging.error(f"Unexpected error: {e}") + raise DatabaseError(f"Unexpected error: {e}") + + +def check_existing_media(url): + db = Database() + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute('SELECT id FROM Media WHERE url = ?', (url,)) + result = cursor.fetchone() + return {'id': result[0]} if result else None + except Exception as e: + logging.error(f"Error checking existing media: {e}") + return None + + +# Modified update_media_content function to create a new version +def update_media_content_with_version(media_id, info_dict, content_input, prompt_input, summary_input, whisper_model): + db = Database() + try: + with db.get_connection() as conn: + cursor = conn.cursor() + + # Create new document version + cursor.execute('SELECT MAX(version) FROM MediaVersion WHERE media_id = ?', (media_id,)) + current_version = cursor.fetchone()[0] or 0 + new_version = current_version + 1 + + # Insert new version + cursor.execute(''' + INSERT INTO MediaVersion (media_id, version, prompt, summary, created_at) + VALUES (?, ?, ?, ?, ?) + ''', (media_id, new_version, prompt_input, summary_input, datetime.now().strftime('%Y-%m-%d %H:%M:%S'))) + + # Update the main content in the Media table + cursor.execute(''' + UPDATE Media + SET content = ?, transcription_model = ?, title = ?, author = ?, ingestion_date = ?, chunking_status = ? + WHERE id = ? + ''', (content_input, whisper_model, info_dict.get('title', 'Untitled'), + info_dict.get('uploader', 'Unknown'), datetime.now().strftime('%Y-%m-%d'), 'pending', media_id)) + + # Update or insert into MediaModifications + cursor.execute(''' + INSERT OR REPLACE INTO MediaModifications (media_id, prompt, summary, modification_date) + VALUES (?, ?, ?, ?) + ''', (media_id, prompt_input, summary_input, datetime.now().strftime('%Y-%m-%d'))) + + # Update full-text search index + cursor.execute('INSERT OR REPLACE INTO media_fts (rowid, title, content) VALUES (?, ?, ?)', + (media_id, info_dict.get('title', 'Untitled'), content_input)) + + conn.commit() + + # Schedule chunking + schedule_chunking(media_id, content_input, info_dict.get('title', 'Untitled')) + + return f"Content updated successfully for media ID: {media_id}. New version: {new_version}" + except Exception as e: + logging.error(f"Error updating media content: {e}") + return f"Error updating content: {str(e)}" + + +# FIXME: This function is not complete and needs to be implemented +def schedule_chunking(media_id: int, content: str, media_name: str): + try: + chunks = chunk_text(content, chunk_options['method'], chunk_options['max_size'], chunk_options['overlap']) + db = Database() + with db.get_connection() as conn: + cursor = conn.cursor() + for i, chunk in enumerate(chunks): + cursor.execute(''' + INSERT INTO MediaChunks (media_id, chunk_text, start_index, end_index, chunk_id) + VALUES (?, ?, ?, ?, ?) + ''', (media_id, chunk, i * chunk_options['max_size'], + min((i + 1) * chunk_options['max_size'], len(content)), + f"{media_id}_chunk_{i}")) + conn.commit() + + # Update chunking status + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute("UPDATE Media SET chunking_status = 'completed' WHERE id = ?", (media_id,)) + conn.commit() + + except Exception as e: + logging.error(f"Error scheduling chunking for media_id {media_id}: {str(e)}") + # You might want to update the chunking_status to 'failed' here + +# +# End of .... +####################################################################################################################### + + +####################################################################################################################### +# +# Functions to manage prompts DB + +def create_prompts_db(): + logging.debug("create_prompts_db: Creating prompts database.") + with sqlite3.connect(get_database_path('prompts.db')) as conn: + cursor = conn.cursor() + cursor.executescript(''' + CREATE TABLE IF NOT EXISTS Prompts ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT NOT NULL UNIQUE, + author TEXT, + details TEXT, + system TEXT, + user TEXT + ); + CREATE TABLE IF NOT EXISTS Keywords ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + keyword TEXT NOT NULL UNIQUE COLLATE NOCASE + ); + CREATE TABLE IF NOT EXISTS PromptKeywords ( + prompt_id INTEGER, + keyword_id INTEGER, + FOREIGN KEY (prompt_id) REFERENCES Prompts (id), + FOREIGN KEY (keyword_id) REFERENCES Keywords (id), + PRIMARY KEY (prompt_id, keyword_id) + ); + CREATE INDEX IF NOT EXISTS idx_keywords_keyword ON Keywords(keyword); + CREATE INDEX IF NOT EXISTS idx_promptkeywords_prompt_id ON PromptKeywords(prompt_id); + CREATE INDEX IF NOT EXISTS idx_promptkeywords_keyword_id ON PromptKeywords(keyword_id); + ''') + +# FIXME - dirty hack that should be removed later... +# Migration function to add the 'author' column to the Prompts table +def add_author_column_to_prompts(): + with sqlite3.connect(get_database_path('prompts.db')) as conn: + cursor = conn.cursor() + # Check if 'author' column already exists + cursor.execute("PRAGMA table_info(Prompts)") + columns = [col[1] for col in cursor.fetchall()] + + if 'author' not in columns: + # Add the 'author' column + cursor.execute('ALTER TABLE Prompts ADD COLUMN author TEXT') + print("Author column added to Prompts table.") + else: + print("Author column already exists in Prompts table.") + +add_author_column_to_prompts() + +def normalize_keyword(keyword): + return re.sub(r'\s+', ' ', keyword.strip().lower()) + + +# FIXME - update calls to this function to use the new args +def add_prompt(name, author, details, system=None, user=None, keywords=None): + logging.debug(f"add_prompt: Adding prompt with name: {name}, author: {author}, system: {system}, user: {user}, keywords: {keywords}") + if not name: + logging.error("add_prompt: A name is required.") + return "A name is required." + + try: + with sqlite3.connect(get_database_path('prompts.db')) as conn: + cursor = conn.cursor() + cursor.execute(''' + INSERT INTO Prompts (name, author, details, system, user) + VALUES (?, ?, ?, ?, ?) + ''', (name, author, details, system, user)) + prompt_id = cursor.lastrowid + + if keywords: + normalized_keywords = [normalize_keyword(k) for k in keywords if k.strip()] + for keyword in set(normalized_keywords): # Use set to remove duplicates + cursor.execute(''' + INSERT OR IGNORE INTO Keywords (keyword) VALUES (?) + ''', (keyword,)) + cursor.execute('SELECT id FROM Keywords WHERE keyword = ?', (keyword,)) + keyword_id = cursor.fetchone()[0] + cursor.execute(''' + INSERT OR IGNORE INTO PromptKeywords (prompt_id, keyword_id) VALUES (?, ?) + ''', (prompt_id, keyword_id)) + return "Prompt added successfully." + except sqlite3.IntegrityError: + return "Prompt with this name already exists." + except sqlite3.Error as e: + return f"Database error: {e}" + + +def fetch_prompt_details(name): + logging.debug(f"fetch_prompt_details: Fetching details for prompt: {name}") + with sqlite3.connect(get_database_path('prompts.db')) as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT p.name, p.author, p.details, p.system, p.user, GROUP_CONCAT(k.keyword, ', ') as keywords + FROM Prompts p + LEFT JOIN PromptKeywords pk ON p.id = pk.prompt_id + LEFT JOIN Keywords k ON pk.keyword_id = k.id + WHERE p.name = ? + GROUP BY p.id + ''', (name,)) + return cursor.fetchone() + + +def list_prompts(page=1, per_page=10): + logging.debug(f"list_prompts: Listing prompts for page {page} with {per_page} prompts per page.") + offset = (page - 1) * per_page + with sqlite3.connect(get_database_path('prompts.db')) as conn: + cursor = conn.cursor() + cursor.execute('SELECT name FROM Prompts LIMIT ? OFFSET ?', (per_page, offset)) + prompts = [row[0] for row in cursor.fetchall()] + + # Get total count of prompts + cursor.execute('SELECT COUNT(*) FROM Prompts') + total_count = cursor.fetchone()[0] + + total_pages = (total_count + per_page - 1) // per_page + return prompts, total_pages, page + +# This will not scale. For a large number of prompts, use a more efficient method. +# FIXME - see above statement. +def load_preset_prompts(): + logging.debug("load_preset_prompts: Loading preset prompts.") + try: + with sqlite3.connect(get_database_path('prompts.db')) as conn: + cursor = conn.cursor() + cursor.execute('SELECT name FROM Prompts ORDER BY name ASC') + prompts = [row[0] for row in cursor.fetchall()] + return prompts + except sqlite3.Error as e: + print(f"Database error: {e}") + return [] + + +def insert_prompt_to_db(title, author, description, system_prompt, user_prompt, keywords=None): + return add_prompt(title, author, description, system_prompt, user_prompt, keywords) + + +def get_prompt_db_connection(): + prompt_db_path = get_database_path('prompts.db') + return sqlite3.connect(prompt_db_path) + + +def search_prompts(query): + logging.debug(f"search_prompts: Searching prompts with query: {query}") + try: + with get_prompt_db_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + SELECT p.name, p.details, p.system, p.user, GROUP_CONCAT(k.keyword, ', ') as keywords + FROM Prompts p + LEFT JOIN PromptKeywords pk ON p.id = pk.prompt_id + LEFT JOIN Keywords k ON pk.keyword_id = k.id + WHERE p.name LIKE ? OR p.details LIKE ? OR p.system LIKE ? OR p.user LIKE ? OR k.keyword LIKE ? + GROUP BY p.id + ORDER BY p.name + """, (f'%{query}%', f'%{query}%', f'%{query}%', f'%{query}%', f'%{query}%')) + return cursor.fetchall() + except sqlite3.Error as e: + logging.error(f"Error searching prompts: {e}") + return [] + + +def search_prompts_by_keyword(keyword, page=1, per_page=10): + logging.debug(f"search_prompts_by_keyword: Searching prompts by keyword: {keyword}") + normalized_keyword = normalize_keyword(keyword) + offset = (page - 1) * per_page + with sqlite3.connect(get_database_path('prompts.db')) as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT DISTINCT p.name + FROM Prompts p + JOIN PromptKeywords pk ON p.id = pk.prompt_id + JOIN Keywords k ON pk.keyword_id = k.id + WHERE k.keyword LIKE ? + LIMIT ? OFFSET ? + ''', ('%' + normalized_keyword + '%', per_page, offset)) + prompts = [row[0] for row in cursor.fetchall()] + + # Get total count of matching prompts + cursor.execute(''' + SELECT COUNT(DISTINCT p.id) + FROM Prompts p + JOIN PromptKeywords pk ON p.id = pk.prompt_id + JOIN Keywords k ON pk.keyword_id = k.id + WHERE k.keyword LIKE ? + ''', ('%' + normalized_keyword + '%',)) + total_count = cursor.fetchone()[0] + + total_pages = (total_count + per_page - 1) // per_page + return prompts, total_pages, page + + +def update_prompt_keywords(prompt_name, new_keywords): + logging.debug(f"update_prompt_keywords: Updating keywords for prompt: {prompt_name}") + try: + with sqlite3.connect(get_database_path('prompts.db')) as conn: + cursor = conn.cursor() + + cursor.execute('SELECT id FROM Prompts WHERE name = ?', (prompt_name,)) + prompt_id = cursor.fetchone() + if not prompt_id: + return "Prompt not found." + prompt_id = prompt_id[0] + + cursor.execute('DELETE FROM PromptKeywords WHERE prompt_id = ?', (prompt_id,)) + + normalized_keywords = [normalize_keyword(k) for k in new_keywords if k.strip()] + for keyword in set(normalized_keywords): # Use set to remove duplicates + cursor.execute('INSERT OR IGNORE INTO Keywords (keyword) VALUES (?)', (keyword,)) + cursor.execute('SELECT id FROM Keywords WHERE keyword = ?', (keyword,)) + keyword_id = cursor.fetchone()[0] + cursor.execute('INSERT INTO PromptKeywords (prompt_id, keyword_id) VALUES (?, ?)', + (prompt_id, keyword_id)) + + # Remove unused keywords + cursor.execute(''' + DELETE FROM Keywords + WHERE id NOT IN (SELECT DISTINCT keyword_id FROM PromptKeywords) + ''') + return "Keywords updated successfully." + except sqlite3.Error as e: + return f"Database error: {e}" + + +def add_or_update_prompt(title, author, description, system_prompt, user_prompt, keywords=None): + logging.debug(f"add_or_update_prompt: Adding or updating prompt: {title}") + if not title: + return "Error: Title is required." + + existing_prompt = fetch_prompt_details(title) + if existing_prompt: + # Update existing prompt + result = update_prompt_in_db(title, author, description, system_prompt, user_prompt) + if "successfully" in result: + # Update keywords if the prompt update was successful + keyword_result = update_prompt_keywords(title, keywords or []) + result += f" {keyword_result}" + else: + # Insert new prompt + result = insert_prompt_to_db(title, author, description, system_prompt, user_prompt, keywords) + + return result + + +def load_prompt_details(selected_prompt): + logging.debug(f"load_prompt_details: Loading prompt details for {selected_prompt}") + if selected_prompt: + details = fetch_prompt_details(selected_prompt) + if details: + return details[0], details[1], details[2], details[3], details[4], details[5] + return "", "", "", "", "", "" + + +def update_prompt_in_db(title, author, description, system_prompt, user_prompt): + logging.debug(f"update_prompt_in_db: Updating prompt: {title}") + try: + with sqlite3.connect(get_database_path('prompts.db')) as conn: + cursor = conn.cursor() + cursor.execute( + "UPDATE Prompts SET author = ?, details = ?, system = ?, user = ? WHERE name = ?", + (author, description, system_prompt, user_prompt, title) + ) + if cursor.rowcount == 0: + return "No prompt found with the given title." + return "Prompt updated successfully!" + except sqlite3.Error as e: + return f"Error updating prompt: {e}" + + +create_prompts_db() + +def delete_prompt(prompt_id): + logging.debug(f"delete_prompt: Deleting prompt with ID: {prompt_id}") + try: + with sqlite3.connect(get_database_path('prompts.db')) as conn: + cursor = conn.cursor() + + # Delete associated keywords + cursor.execute("DELETE FROM PromptKeywords WHERE prompt_id = ?", (prompt_id,)) + + # Delete the prompt + cursor.execute("DELETE FROM Prompts WHERE id = ?", (prompt_id,)) + + if cursor.rowcount == 0: + return f"No prompt found with ID {prompt_id}" + else: + conn.commit() + return f"Prompt with ID {prompt_id} has been successfully deleted" + except sqlite3.Error as e: + return f"An error occurred: {e}" + +# +# +####################################################################################################################### + + +####################################################################################################################### +# +# Function to fetch/update media content + +def update_media_content(selected_item, item_mapping, content_input, prompt_input, summary_input): + try: + if selected_item and item_mapping and selected_item in item_mapping: + media_id = item_mapping[selected_item] + + with db.get_connection() as conn: + cursor = conn.cursor() + + # Update the main content in the Media table + cursor.execute("UPDATE Media SET content = ? WHERE id = ?", (content_input, media_id)) + + # Check if a row already exists in MediaModifications for this media_id + cursor.execute("SELECT COUNT(*) FROM MediaModifications WHERE media_id = ?", (media_id,)) + exists = cursor.fetchone()[0] > 0 + + if exists: + # Update existing row + cursor.execute(""" + UPDATE MediaModifications + SET prompt = ?, summary = ?, modification_date = CURRENT_TIMESTAMP + WHERE media_id = ? + """, (prompt_input, summary_input, media_id)) + else: + # Insert new row + cursor.execute(""" + INSERT INTO MediaModifications (media_id, prompt, summary, modification_date) + VALUES (?, ?, ?, CURRENT_TIMESTAMP) + """, (media_id, prompt_input, summary_input)) + + # Create new document version + new_version = create_document_version(media_id, content_input) + + conn.commit() + + return f"Content updated successfully for media ID: {media_id}. New version: {new_version}" + else: + return "No item selected or invalid selection" + except Exception as e: + logging.error(f"Error updating media content: {e}") + return f"Error updating content: {str(e)}" + + +def search_media_database(query: str, connection=None) -> List[Tuple[int, str, str]]: + def execute_query(conn): + try: + cursor = conn.cursor() + cursor.execute("SELECT id, title, url FROM Media WHERE title LIKE ?", (f'%{query}%',)) + return cursor.fetchall() + except sqlite3.Error as e: + raise Exception(f"Error searching media database: {e}") + + if connection: + return execute_query(connection) + else: + with db.get_connection() as conn: + return execute_query(conn) + + +def load_media_content(media_id: int) -> dict: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute("SELECT content, prompt, summary FROM Media WHERE id = ?", (media_id,)) + result = cursor.fetchone() + if result: + return { + "content": result[0], + "prompt": result[1], + "summary": result[2] + } + return {"content": "", "prompt": "", "summary": ""} + except sqlite3.Error as e: + raise Exception(f"Error loading media content: {e}") + + +def fetch_items_by_title_or_url(search_query: str, search_type: str): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + if search_type == 'Title': + cursor.execute("SELECT id, title, url FROM Media WHERE title LIKE ?", (f'%{search_query}%',)) + elif search_type == 'URL': + cursor.execute("SELECT id, title, url FROM Media WHERE url LIKE ?", (f'%{search_query}%',)) + results = cursor.fetchall() + return results + except sqlite3.Error as e: + raise DatabaseError(f"Error fetching items by {search_type}: {e}") + + +def fetch_items_by_keyword(search_query: str): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + SELECT m.id, m.title, m.url + FROM Media m + JOIN MediaKeywords mk ON m.id = mk.media_id + JOIN Keywords k ON mk.keyword_id = k.id + WHERE k.keyword LIKE ? + """, (f'%{search_query}%',)) + results = cursor.fetchall() + return results + except sqlite3.Error as e: + raise DatabaseError(f"Error fetching items by keyword: {e}") + + +def fetch_items_by_content(search_query: str): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute("SELECT id, title, url FROM Media WHERE content LIKE ?", (f'%{search_query}%',)) + results = cursor.fetchall() + return results + except sqlite3.Error as e: + raise DatabaseError(f"Error fetching items by content: {e}") + + +def fetch_item_details_single(media_id: int): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + SELECT prompt, summary + FROM MediaModifications + WHERE media_id = ? + ORDER BY modification_date DESC + LIMIT 1 + """, (media_id,)) + prompt_summary_result = cursor.fetchone() + cursor.execute("SELECT content FROM Media WHERE id = ?", (media_id,)) + content_result = cursor.fetchone() + + prompt = prompt_summary_result[0] if prompt_summary_result else "No prompt available." + summary = prompt_summary_result[1] if prompt_summary_result else "No summary available." + content = content_result[0] if content_result else "No content available." + + return prompt, summary, content + except sqlite3.Error as e: + logging.error(f"Error fetching item details: {e}") + return "Error fetching prompt.", "Error fetching summary.", "Error fetching content." + + + +def convert_to_markdown(item): + markdown = f"# {item['title']}\n\n" + markdown += f"**URL:** {item['url']}\n\n" + markdown += f"**Author:** {item['author']}\n\n" + markdown += f"**Ingestion Date:** {item['ingestion_date']}\n\n" + markdown += f"**Type:** {item['type']}\n\n" + markdown += f"**Keywords:** {', '.join(item['keywords'])}\n\n" + markdown += "## Summary\n\n" + markdown += f"{item['summary']}\n\n" + markdown += "## Content\n\n" + markdown += f"{item['content']}\n\n" + return markdown + +# Gradio function to handle user input and display results with pagination for displaying entries in the DB +def fetch_paginated_data(page: int, results_per_page: int) -> Tuple[List[Tuple], int]: + try: + offset = (page - 1) * results_per_page + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute("SELECT COUNT(*) FROM Media") + total_entries = cursor.fetchone()[0] + + cursor.execute("SELECT id, title, url FROM Media LIMIT ? OFFSET ?", (results_per_page, offset)) + results = cursor.fetchall() + + return results, total_entries + except sqlite3.Error as e: + raise Exception(f"Error fetching paginated data: {e}") + +def format_results_as_html(results: List[Tuple]) -> str: + html = "" + html += "" + for row in results: + html += f"" + html += "
IDTitleURL
{row[0]}{row[1]}{row[2]}
" + return html + +def view_database(page: int, results_per_page: int) -> Tuple[str, str, int]: + results, total_entries = fetch_paginated_data(page, results_per_page) + formatted_results = format_results_as_html(results) + # Calculate total pages + total_pages = (total_entries + results_per_page - 1) // results_per_page + return formatted_results, f"Page {page} of {total_pages}", total_pages + + +def search_and_display_items(query, search_type, page, entries_per_page,char_count): + offset = (page - 1) * entries_per_page + try: + with sqlite3.connect('./Databases/media_summary.db') as conn: + cursor = conn.cursor() + + # Adjust the SQL query based on the search type + if search_type == "Title": + where_clause = "WHERE m.title LIKE ?" + elif search_type == "URL": + where_clause = "WHERE m.url LIKE ?" + elif search_type == "Keyword": + where_clause = "WHERE k.keyword LIKE ?" + elif search_type == "Content": + where_clause = "WHERE m.content LIKE ?" + else: + raise ValueError("Invalid search type") + + cursor.execute(f''' + SELECT m.id, m.title, m.url, m.content, mm.summary, GROUP_CONCAT(k.keyword, ', ') as keywords + FROM Media m + LEFT JOIN MediaModifications mm ON m.id = mm.media_id + LEFT JOIN MediaKeywords mk ON m.id = mk.media_id + LEFT JOIN Keywords k ON mk.keyword_id = k.id + {where_clause} + GROUP BY m.id + ORDER BY m.ingestion_date DESC + LIMIT ? OFFSET ? + ''', (f'%{query}%', entries_per_page, offset)) + items = cursor.fetchall() + + cursor.execute(f''' + SELECT COUNT(DISTINCT m.id) + FROM Media m + LEFT JOIN MediaKeywords mk ON m.id = mk.media_id + LEFT JOIN Keywords k ON mk.keyword_id = k.id + {where_clause} + ''', (f'%{query}%',)) + total_items = cursor.fetchone()[0] + + results = "" + for item in items: + title = html.escape(item[1]).replace('\n', '
') + url = html.escape(item[2]).replace('\n', '
') + # First X amount of characters of the content + content = html.escape(item[3] or '')[:char_count] + '...' + summary = html.escape(item[4] or '').replace('\n', '
') + keywords = html.escape(item[5] or '').replace('\n', '
') + + results += f""" +
+
+
Title: {title}
+
URL: {url}
+
+
+ Content (first {char_count} characters): + 
{content}
+
+
+ Summary: + 
{summary}
+
+
+ Keywords: {keywords} +
+
+ """ + + total_pages = (total_items + entries_per_page - 1) // entries_per_page + pagination = f"Page {page} of {total_pages} (Total items: {total_items})" + + return results, pagination, total_pages + except sqlite3.Error as e: + return f"

Error searching items: {e}

", "Error", 0 + + +# +# End of Functions to manage prompts DB / Fetch and update media content +####################################################################################################################### + + +####################################################################################################################### +# +# Obsidian-related Functions + +def import_obsidian_note_to_db(note_data): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + + cursor.execute("SELECT id FROM Media WHERE title = ? AND type = 'obsidian_note'", (note_data['title'],)) + existing_note = cursor.fetchone() + + # Generate a relative path or meaningful identifier instead of using the temporary file path + relative_path = os.path.relpath(note_data['file_path'], start=os.path.dirname(note_data['file_path'])) + + if existing_note: + media_id = existing_note[0] + cursor.execute(""" + UPDATE Media + SET content = ?, author = ?, ingestion_date = CURRENT_TIMESTAMP, url = ? + WHERE id = ? + """, (note_data['content'], note_data['frontmatter'].get('author', 'Unknown'), relative_path, media_id)) + + cursor.execute("DELETE FROM MediaKeywords WHERE media_id = ?", (media_id,)) + else: + cursor.execute(""" + INSERT INTO Media (title, content, type, author, ingestion_date, url) + VALUES (?, ?, 'obsidian_note', ?, CURRENT_TIMESTAMP, ?) + """, (note_data['title'], note_data['content'], note_data['frontmatter'].get('author', 'Unknown'), + relative_path)) + + media_id = cursor.lastrowid + + for tag in note_data['tags']: + cursor.execute("INSERT OR IGNORE INTO Keywords (keyword) VALUES (?)", (tag,)) + cursor.execute("SELECT id FROM Keywords WHERE keyword = ?", (tag,)) + keyword_id = cursor.fetchone()[0] + cursor.execute("INSERT OR IGNORE INTO MediaKeywords (media_id, keyword_id) VALUES (?, ?)", + (media_id, keyword_id)) + + frontmatter_str = yaml.dump(note_data['frontmatter']) + cursor.execute(""" + INSERT INTO MediaModifications (media_id, prompt, summary, modification_date) + VALUES (?, 'Obsidian Frontmatter', ?, CURRENT_TIMESTAMP) + """, (media_id, frontmatter_str)) + + # Update full-text search index + cursor.execute('INSERT OR REPLACE INTO media_fts (rowid, title, content) VALUES (?, ?, ?)', + (media_id, note_data['title'], note_data['content'])) + + action = "Updated" if existing_note else "Imported" + logger.info(f"{action} Obsidian note: {note_data['title']}") + return True, None + except sqlite3.Error as e: + error_msg = f"Database error {'updating' if existing_note else 'importing'} note {note_data['title']}: {str(e)}" + logger.error(error_msg) + return False, error_msg + except Exception as e: + error_msg = f"Unexpected error {'updating' if existing_note else 'importing'} note {note_data['title']}: {str(e)}\n{traceback.format_exc()}" + logger.error(error_msg) + return False, error_msg + + +# +# End of Obsidian-related Functions +####################################################################################################################### + + +####################################################################################################################### +# +# Chat-related Functions + + + +def create_chat_conversation(media_id, conversation_name): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + INSERT INTO ChatConversations (media_id, conversation_name, created_at, updated_at) + VALUES (?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + ''', (media_id, conversation_name)) + conn.commit() + return cursor.lastrowid + except sqlite3.Error as e: + logging.error(f"Error creating chat conversation: {e}") + raise DatabaseError(f"Error creating chat conversation: {e}") + + +def add_chat_message(conversation_id: int, sender: str, message: str) -> int: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + INSERT INTO ChatMessages (conversation_id, sender, message) + VALUES (?, ?, ?) + ''', (conversation_id, sender, message)) + conn.commit() + return cursor.lastrowid + except sqlite3.Error as e: + logging.error(f"Error adding chat message: {e}") + raise DatabaseError(f"Error adding chat message: {e}") + + +def get_chat_messages(conversation_id: int) -> List[Dict[str, Any]]: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT id, sender, message, timestamp + FROM ChatMessages + WHERE conversation_id = ? + ORDER BY timestamp ASC + ''', (conversation_id,)) + messages = cursor.fetchall() + return [ + { + 'id': msg[0], + 'sender': msg[1], + 'message': msg[2], + 'timestamp': msg[3] + } + for msg in messages + ] + except sqlite3.Error as e: + logging.error(f"Error retrieving chat messages: {e}") + raise DatabaseError(f"Error retrieving chat messages: {e}") + + +def search_chat_conversations(search_query: str) -> List[Dict[str, Any]]: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT cc.id, cc.media_id, cc.conversation_name, cc.created_at, m.title as media_title + FROM ChatConversations cc + LEFT JOIN Media m ON cc.media_id = m.id + WHERE cc.conversation_name LIKE ? OR m.title LIKE ? + ORDER BY cc.updated_at DESC + ''', (f'%{search_query}%', f'%{search_query}%')) + conversations = cursor.fetchall() + return [ + { + 'id': conv[0], + 'media_id': conv[1], + 'conversation_name': conv[2], + 'created_at': conv[3], + 'media_title': conv[4] or "Unknown Media" + } + for conv in conversations + ] + except sqlite3.Error as e: + logging.error(f"Error searching chat conversations: {e}") + return [] + + +def update_chat_message(message_id: int, new_message: str) -> None: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + UPDATE ChatMessages + SET message = ?, timestamp = CURRENT_TIMESTAMP + WHERE id = ? + ''', (new_message, message_id)) + conn.commit() + except sqlite3.Error as e: + logging.error(f"Error updating chat message: {e}") + raise DatabaseError(f"Error updating chat message: {e}") + + +def delete_chat_message(message_id: int) -> None: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute('DELETE FROM ChatMessages WHERE id = ?', (message_id,)) + conn.commit() + except sqlite3.Error as e: + logging.error(f"Error deleting chat message: {e}") + raise DatabaseError(f"Error deleting chat message: {e}") + + +def save_chat_history_to_database(chatbot, conversation_id, media_id, media_name, conversation_name): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + + # If conversation_id is None, create a new conversation + if conversation_id is None: + cursor.execute(''' + INSERT INTO ChatConversations (media_id, media_name, conversation_name, created_at, updated_at) + VALUES (?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + ''', (media_id, media_name, conversation_name)) + conversation_id = cursor.lastrowid + else: + # If conversation exists, update the media_name + cursor.execute(''' + UPDATE ChatConversations + SET media_name = ?, updated_at = CURRENT_TIMESTAMP + WHERE id = ? + ''', (media_name, conversation_id)) + + # Save each message in the chatbot history + for i, (user_msg, ai_msg) in enumerate(chatbot): + cursor.execute(''' + INSERT INTO ChatMessages (conversation_id, sender, message, timestamp) + VALUES (?, ?, ?, CURRENT_TIMESTAMP) + ''', (conversation_id, 'user', user_msg)) + + cursor.execute(''' + INSERT INTO ChatMessages (conversation_id, sender, message, timestamp) + VALUES (?, ?, ?, CURRENT_TIMESTAMP) + ''', (conversation_id, 'ai', ai_msg)) + + # Update the conversation's updated_at timestamp + cursor.execute(''' + UPDATE ChatConversations + SET updated_at = CURRENT_TIMESTAMP + WHERE id = ? + ''', (conversation_id,)) + + conn.commit() + + return conversation_id + except Exception as e: + logging.error(f"Error saving chat history to database: {str(e)}") + raise + + +def get_conversation_name(conversation_id): + if conversation_id is None: + return None + + try: + with sqlite3.connect('media_summary.db') as conn: # Replace with your actual database name + cursor = conn.cursor() + + query = """ + SELECT conversation_name, media_name + FROM ChatConversations + WHERE id = ? + """ + + cursor.execute(query, (conversation_id,)) + result = cursor.fetchone() + + if result: + conversation_name, media_name = result + if conversation_name: + return conversation_name + elif media_name: + return f"{media_name}-chat" + + return None # Return None if no result found + except sqlite3.Error as e: + logging.error(f"Database error in get_conversation_name: {e}") + return None + except Exception as e: + logging.error(f"Unexpected error in get_conversation_name: {e}") + return None + +# +# End of Chat-related Functions +####################################################################################################################### + + +####################################################################################################################### +# +# Functions to Compare Transcripts + +# Fetch Transcripts +def get_transcripts(media_id): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT id, whisper_model, transcription, created_at + FROM Transcripts + WHERE media_id = ? + ORDER BY created_at DESC + ''', (media_id,)) + return cursor.fetchall() + except Exception as e: + logging.error(f"Error in get_transcripts: {str(e)}") + return [] + +def get_latest_transcription(media_id: int): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + SELECT transcription + FROM Transcripts + WHERE media_id = ? + ORDER BY created_at DESC + LIMIT 1 + """, (media_id,)) + result = cursor.fetchone() + return result[0] if result else "No transcription available." + except sqlite3.Error as e: + logging.error(f"Error fetching latest transcription: {e}") + return "Error fetching transcription." + +# +# End of Functions to Compare Transcripts +####################################################################################################################### + + +####################################################################################################################### +# +# Functions to handle deletion of media items + + +def mark_as_trash(media_id: int) -> None: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + UPDATE Media + SET is_trash = 1, trash_date = ? + WHERE id = ? + """, (datetime.now(), media_id)) + conn.commit() + + +def restore_from_trash(media_id: int) -> None: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + UPDATE Media + SET is_trash = 0, trash_date = NULL + WHERE id = ? + """, (media_id,)) + conn.commit() + + +def get_trashed_items() -> List[Dict]: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + SELECT id, title, trash_date + FROM Media + WHERE is_trash = 1 + ORDER BY trash_date DESC + """) + return [{'id': row[0], 'title': row[1], 'trash_date': row[2]} for row in cursor.fetchall()] + + +def permanently_delete_item(media_id: int) -> None: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute("DELETE FROM Media WHERE id = ?", (media_id,)) + cursor.execute("DELETE FROM MediaKeywords WHERE media_id = ?", (media_id,)) + cursor.execute("DELETE FROM MediaVersion WHERE media_id = ?", (media_id,)) + cursor.execute("DELETE FROM MediaModifications WHERE media_id = ?", (media_id,)) + cursor.execute("DELETE FROM media_fts WHERE rowid = ?", (media_id,)) + conn.commit() + + +def empty_trash(days_threshold: int) -> Tuple[int, int]: + threshold_date = datetime.now() - timedelta(days=days_threshold) + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + SELECT id FROM Media + WHERE is_trash = 1 AND trash_date <= ? + """, (threshold_date,)) + old_items = cursor.fetchall() + + for item in old_items: + permanently_delete_item(item[0]) + + cursor.execute(""" + SELECT COUNT(*) FROM Media + WHERE is_trash = 1 AND trash_date > ? + """, (threshold_date,)) + remaining_items = cursor.fetchone()[0] + + return len(old_items), remaining_items + + +def user_delete_item(media_id: int, force: bool = False) -> str: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute("SELECT is_trash, trash_date FROM Media WHERE id = ?", (media_id,)) + result = cursor.fetchone() + + if not result: + return "Item not found." + + is_trash, trash_date = result + + if not is_trash: + mark_as_trash(media_id) + return "Item moved to trash." + + if force or (trash_date and (datetime.now() - trash_date).days >= 30): + permanently_delete_item(media_id) + return "Item permanently deleted." + else: + return "Item is already in trash. Use force=True to delete permanently before 30 days." + + +def get_chunk_text(media_id: int, chunk_index: int) -> str: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute("SELECT content FROM MediaChunks WHERE media_id = ? AND chunk_index = ?", + (media_id, chunk_index)) + result = cursor.fetchone() + return result[0] if result else None + +def get_full_document(media_id: int) -> str: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute("SELECT content FROM Media WHERE id = ?", (media_id,)) + result = cursor.fetchone() + return result[0] if result else None + +def get_all_content_from_database() -> List[Dict[str, Any]]: + """ + Retrieve all media content from the database that requires embedding. + + Returns: + List[Dict[str, Any]]: A list of dictionaries, each containing the media ID, content, title, and other relevant fields. + """ + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + SELECT id, content, title, author, type + FROM Media + WHERE is_trash = 0 -- Exclude items marked as trash + """) + media_items = cursor.fetchall() + + all_content = [ + { + 'id': item[0], + 'content': item[1], + 'title': item[2], + 'author': item[3], + 'type': item[4] + } + for item in media_items + ] + + return all_content + + except sqlite3.Error as e: + logger.error(f"Error retrieving all content from database: {e}") + raise DatabaseError(f"Error retrieving all content from database: {e}") + + +def get_media_content(media_id: int) -> str: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute("SELECT content FROM Media WHERE id = ?", (media_id,)) + result = cursor.fetchone() + if result is None: + raise ValueError(f"No media found with id {media_id}") + return result[0] + except sqlite3.Error as e: + logging.error(f"Database error in get_media_content: {e}") + raise DatabaseError(f"Failed to retrieve media content: {e}") + except Exception as e: + logging.error(f"Unexpected error in get_media_content: {e}") + raise + +def get_media_title(media_id: int) -> str: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute("SELECT title FROM Media WHERE id = ?", (media_id,)) + result = cursor.fetchone() + return result[0] if result else f"Unknown Source (ID: {media_id})" + except sqlite3.Error as e: + logging.error(f"Database error in get_media_title: {e}") + return f"Unknown Source (ID: {media_id})" + +def get_media_transcripts(media_id): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT id, whisper_model, transcription, created_at + FROM Transcripts + WHERE media_id = ? + ORDER BY created_at DESC + ''', (media_id,)) + results = cursor.fetchall() + return [ + { + 'id': row[0], + 'whisper_model': row[1], + 'content': row[2], + 'created_at': row[3] + } + for row in results + ] + except Exception as e: + logging.error(f"Error in get_media_transcripts: {str(e)}") + return [] + +def get_specific_transcript(transcript_id: int) -> Dict: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT id, whisper_model, transcription, created_at + FROM Transcripts + WHERE id = ? + ''', (transcript_id,)) + result = cursor.fetchone() + if result: + return { + 'id': result[0], + 'whisper_model': result[1], + 'content': result[2], + 'created_at': result[3] + } + return {'error': f"No transcript found with ID {transcript_id}"} + except Exception as e: + logging.error(f"Error in get_specific_transcript: {str(e)}") + return {'error': f"Error retrieving transcript: {str(e)}"} + +def get_media_summaries(media_id: int) -> List[Dict]: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT id, summary, modification_date + FROM MediaModifications + WHERE media_id = ? AND summary IS NOT NULL + ORDER BY modification_date DESC + ''', (media_id,)) + results = cursor.fetchall() + return [ + { + 'id': row[0], + 'content': row[1], + 'created_at': row[2] + } + for row in results + ] + except Exception as e: + logging.error(f"Error in get_media_summaries: {str(e)}") + +def get_specific_summary(summary_id: int) -> Dict: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT id, summary, modification_date + FROM MediaModifications + WHERE id = ? + ''', (summary_id,)) + result = cursor.fetchone() + if result: + return { + 'id': result[0], + 'content': result[1], + 'created_at': result[2] + } + return {'error': f"No summary found with ID {summary_id}"} + except Exception as e: + logging.error(f"Error in get_specific_summary: {str(e)}") + return {'error': f"Error retrieving summary: {str(e)}"} + +def get_media_prompts(media_id: int) -> List[Dict]: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT id, prompt, modification_date + FROM MediaModifications + WHERE media_id = ? AND prompt IS NOT NULL + ORDER BY modification_date DESC + ''', (media_id,)) + results = cursor.fetchall() + return [ + { + 'id': row[0], + 'content': row[1], + 'created_at': row[2] + } + for row in results + ] + except Exception as e: + logging.error(f"Error in get_media_prompts: {str(e)}") + return [] + +def get_specific_prompt(prompt_id: int) -> Dict: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT id, prompt, modification_date + FROM MediaModifications + WHERE id = ? + ''', (prompt_id,)) + result = cursor.fetchone() + if result: + return { + 'id': result[0], + 'content': result[1], + 'created_at': result[2] + } + return {'error': f"No prompt found with ID {prompt_id}"} + except Exception as e: + logging.error(f"Error in get_specific_prompt: {str(e)}") + return {'error': f"Error retrieving prompt: {str(e)}"} + + +def delete_specific_transcript(transcript_id: int) -> str: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute('DELETE FROM Transcripts WHERE id = ?', (transcript_id,)) + conn.commit() + if cursor.rowcount > 0: + return f"Transcript with ID {transcript_id} has been deleted successfully." + else: + return f"No transcript found with ID {transcript_id}." + except Exception as e: + logging.error(f"Error in delete_specific_transcript: {str(e)}") + return f"Error deleting transcript: {str(e)}" + +def delete_specific_summary(summary_id: int) -> str: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute('UPDATE MediaModifications SET summary = NULL WHERE id = ?', (summary_id,)) + conn.commit() + if cursor.rowcount > 0: + return f"Summary with ID {summary_id} has been deleted successfully." + else: + return f"No summary found with ID {summary_id}." + except Exception as e: + logging.error(f"Error in delete_specific_summary: {str(e)}") + return f"Error deleting summary: {str(e)}" + +def delete_specific_prompt(prompt_id: int) -> str: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute('UPDATE MediaModifications SET prompt = NULL WHERE id = ?', (prompt_id,)) + conn.commit() + if cursor.rowcount > 0: + return f"Prompt with ID {prompt_id} has been deleted successfully." + else: + return f"No prompt found with ID {prompt_id}." + except Exception as e: + logging.error(f"Error in delete_specific_prompt: {str(e)}") + return f"Error deleting prompt: {str(e)}" + + +def get_paginated_files(page: int = 1, results_per_page: int = 50) -> Tuple[List[Tuple[int, str]], int, int]: + try: + offset = (page - 1) * results_per_page + with db.get_connection() as conn: + cursor = conn.cursor() + + # Get total count of media items + cursor.execute("SELECT COUNT(*) FROM Media") + total_entries = cursor.fetchone()[0] + + # Fetch paginated results + cursor.execute(""" + SELECT id, title + FROM Media + ORDER BY title + LIMIT ? OFFSET ? + """, (results_per_page, offset)) + results = cursor.fetchall() + + # Calculate total pages + total_pages = (total_entries + results_per_page - 1) // results_per_page + + return results, total_pages, page + except sqlite3.Error as e: + logging.error(f"Error fetching paginated files: {e}") + raise DatabaseError(f"Error fetching paginated files: {e}") + + +# +# End of Functions to handle deletion of media items +####################################################################################################################### + + +####################################################################################################################### +# +# Functions to manage document versions + +def create_document_version(media_id: int, content: str) -> int: + logging.info(f"Attempting to create document version for media_id: {media_id}") + try: + with db.get_connection() as conn: + cursor = conn.cursor() + + # Start a transaction + cursor.execute("BEGIN EXCLUSIVE TRANSACTION") + + try: + # Verify media_id exists and get the latest version in one query + cursor.execute(''' + SELECT m.id, COALESCE(MAX(dv.version_number), 0) + FROM Media m + LEFT JOIN DocumentVersions dv ON m.id = dv.media_id + WHERE m.id = ? + GROUP BY m.id + ''', (media_id,)) + result = cursor.fetchone() + + if not result: + raise ValueError(f"No Media entry found for id: {media_id}") + + _, latest_version = result + new_version = latest_version + 1 + + logging.debug(f"Inserting new version {new_version} for media_id: {media_id}") + + # Insert new version + cursor.execute(''' + INSERT INTO DocumentVersions (media_id, version_number, content) + VALUES (?, ?, ?) + ''', (media_id, new_version, content)) + + # Commit the transaction + conn.commit() + logging.info(f"Successfully created document version {new_version} for media_id: {media_id}") + return new_version + except Exception as e: + # If any error occurs, roll back the transaction + conn.rollback() + raise e + except sqlite3.Error as e: + logging.error(f"Database error creating document version: {e}") + logging.error(f"Error details - media_id: {media_id}, content length: {len(content)}") + raise DatabaseError(f"Failed to create document version: {e}") + except Exception as e: + logging.error(f"Unexpected error creating document version: {e}") + logging.error(f"Error details - media_id: {media_id}, content length: {len(content)}") + raise + + +def get_document_version(media_id: int, version_number: int = None) -> Dict[str, Any]: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + + if version_number is None: + # Get the latest version + cursor.execute(''' + SELECT id, version_number, content, created_at + FROM DocumentVersions + WHERE media_id = ? + ORDER BY version_number DESC + LIMIT 1 + ''', (media_id,)) + else: + cursor.execute(''' + SELECT id, version_number, content, created_at + FROM DocumentVersions + WHERE media_id = ? AND version_number = ? + ''', (media_id, version_number)) + + result = cursor.fetchone() + + if result: + return { + 'id': result[0], + 'version_number': result[1], + 'content': result[2], + 'created_at': result[3] + } + else: + return {'error': f"No document version found for media_id {media_id}" + (f" and version_number {version_number}" if version_number is not None else "")} + except sqlite3.Error as e: + error_message = f"Error retrieving document version: {e}" + logging.error(error_message) + return {'error': error_message} + +def get_all_document_versions(media_id: int) -> List[Dict[str, Any]]: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + SELECT id, version_number, content, created_at + FROM DocumentVersions + WHERE media_id = ? + ORDER BY version_number DESC + ''', (media_id,)) + results = cursor.fetchall() + + if results: + return [ + { + 'id': row[0], + 'version_number': row[1], + 'content': row[2], + 'created_at': row[3] + } + for row in results + ] + else: + return [] + except sqlite3.Error as e: + error_message = f"Error retrieving all document versions: {e}" + logging.error(error_message) + return [{'error': error_message}] + +def delete_document_version(media_id: int, version_number: int) -> Dict[str, Any]: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + DELETE FROM DocumentVersions + WHERE media_id = ? AND version_number = ? + ''', (media_id, version_number)) + conn.commit() + + if cursor.rowcount > 0: + return {'success': f"Document version {version_number} for media_id {media_id} deleted successfully"} + else: + return {'error': f"No document version found for media_id {media_id} and version_number {version_number}"} + except sqlite3.Error as e: + error_message = f"Error deleting document version: {e}" + logging.error(error_message) + return {'error': error_message} + +def rollback_to_version(media_id: int, version_number: int) -> Dict[str, Any]: + try: + with db.get_connection() as conn: + cursor = conn.cursor() + + # Get the content of the version to rollback to + cursor.execute(''' + SELECT content + FROM DocumentVersions + WHERE media_id = ? AND version_number = ? + ''', (media_id, version_number)) + result = cursor.fetchone() + + if not result: + return {'error': f"No document version found for media_id {media_id} and version_number {version_number}"} + + rollback_content = result[0] + + # Create a new version with the content of the version to rollback to + cursor.execute(''' + INSERT INTO DocumentVersions (media_id, version_number, content) + VALUES (?, (SELECT COALESCE(MAX(version_number), 0) + 1 FROM DocumentVersions WHERE media_id = ?), ?) + ''', (media_id, media_id, rollback_content)) + + new_version_number = cursor.lastrowid + + conn.commit() + + return { + 'success': f"Rolled back to version {version_number} for media_id {media_id}", + 'new_version_number': new_version_number + } + except sqlite3.Error as e: + error_message = f"Error rolling back to document version: {e}" + logging.error(error_message) + return {'error': error_message} + +# +# End of Functions to manage document versions +####################################################################################################################### + + +####################################################################################################################### +# +# Functions to manage media chunks + +def process_chunks(database, chunks: List[Dict], media_id: int, batch_size: int = 100): + """ + Process chunks in batches and insert them into the database. + + :param database: Database instance to use for inserting chunks + :param chunks: List of chunk dictionaries + :param media_id: ID of the media these chunks belong to + :param batch_size: Number of chunks to process in each batch + """ + total_chunks = len(chunks) + processed_chunks = 0 + + for i in range(0, total_chunks, batch_size): + batch = chunks[i:i + batch_size] + chunk_data = [ + (media_id, chunk['text'], chunk['start_index'], chunk['end_index']) + for chunk in batch + ] + + try: + database.execute_many( + "INSERT INTO MediaChunks (media_id, chunk_text, start_index, end_index) VALUES (?, ?, ?, ?)", + chunk_data + ) + processed_chunks += len(batch) + logging.info(f"Processed {processed_chunks}/{total_chunks} chunks for media_id {media_id}") + except Exception as e: + logging.error(f"Error inserting chunk batch for media_id {media_id}: {e}") + # Optionally, you could raise an exception here to stop processing + # raise + + logging.info(f"Finished processing all {total_chunks} chunks for media_id {media_id}") + + +# Usage example: +# chunks = [{'text': 'chunk1', 'start_index': 0, 'end_index': 10}, ...] +# process_chunks(db, chunks, media_id=1, batch_size=100) + +def batch_insert_chunks(conn, chunks, media_id): + cursor = conn.cursor() + chunk_data = [( + media_id, + chunk['text'], + chunk['metadata']['start_index'], + chunk['metadata']['end_index'], + f"{media_id}_chunk_{i}" + ) for i, chunk in enumerate(chunks, 1)] + + cursor.executemany(''' + INSERT INTO MediaChunks (media_id, chunk_text, start_index, end_index, chunk_id) + VALUES (?, ?, ?, ?, ?) + ''', chunk_data) + + +chunk_queue = queue.Queue() + +def chunk_processor(): + while True: + chunk_batch = chunk_queue.get() + if chunk_batch is None: + break + try: + with db.get_connection() as conn: + conn.execute("BEGIN TRANSACTION") + try: + batch_insert_chunks(conn, chunk_batch['chunks'], chunk_batch['media_id']) + conn.commit() + except Exception as e: + conn.rollback() + logging.error(f"Error in batch insert: {str(e)}") + except Exception as e: + logging.error(f"Error processing chunk batch: {str(e)}") + finally: + chunk_queue.task_done() + +# Start the chunk processor thread +#chunk_processor_thread = threading.Thread(target=chunk_processor) +#chunk_processor_thread.start() + +# Make sure to properly shut down the chunk processor when your application exits +# def shutdown_chunk_processor(): +# chunk_queue.put(None) +# chunk_processor_thread.join() + +#FIXME - add into main db creation code +def update_media_chunks_table(): + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(''' + CREATE TABLE IF NOT EXISTS MediaChunks_new ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + media_id INTEGER, + chunk_text TEXT, + start_index INTEGER, + end_index INTEGER, + chunk_id TEXT, + FOREIGN KEY (media_id) REFERENCES Media(id) + ) + ''') + cursor.execute(''' + INSERT INTO MediaChunks_new (media_id, chunk_text, start_index, end_index) + SELECT media_id, chunk_text, start_index, end_index FROM MediaChunks + ''') + cursor.execute('DROP TABLE MediaChunks') + cursor.execute('ALTER TABLE MediaChunks_new RENAME TO MediaChunks') + + logger.info("Updated MediaChunks table schema") + +update_media_chunks_table() +# Above function is a dirty hack that should be merged into the initial DB creation statement. This is a placeholder +# FIXME + + +# This is backwards compatibility for older setups. +# Function to add a missing column to the Media table +def add_missing_column_if_not_exists(db, table_name, column_name, column_definition): + try: + # Check if the column already exists in the table + cursor = db.cursor() + cursor.execute(f"PRAGMA table_info({table_name})") + columns = [column[1] for column in cursor.fetchall()] + + # If the column is not found, add it + if column_name not in columns: + logging.info(f"Adding missing column '{column_name}' to table '{table_name}'") + cursor.execute(f"ALTER TABLE {table_name} ADD COLUMN {column_definition}") + db.commit() + logging.info(f"Column '{column_name}' added successfully.") + else: + logging.info(f"Column '{column_name}' already exists in table '{table_name}'") + + except sqlite3.Error as e: + logging.error(f"Error checking or adding column '{column_name}' in table '{table_name}': {e}") + raise + +# Example usage of the function +def update_media_table(db): + # Add chunking_status column if it doesn't exist + add_missing_column_if_not_exists(db, 'Media', 'chunking_status', "TEXT DEFAULT 'pending'") + +# DEADCODE +# # Vector check FIXME/Delete later +# def alter_media_table(db): +# alter_query = ''' +# ALTER TABLE Media ADD COLUMN vector_processing INTEGER DEFAULT 0 +# ''' +# try: +# db.execute_query(alter_query) +# logging.info("Media table altered successfully to include vector_processing column.") +# except Exception as e: +# logging.error(f"Error altering Media table: {str(e)}") +# # If the column already exists, SQLite will throw an error, which we can safely ignore +# if "duplicate column name" not in str(e).lower(): +# raise +# +# # Vector check FIXME/Delete later +# alter_media_table(db) + +# +# End of Functions to manage media chunks +####################################################################################################################### + + +####################################################################################################################### +# +# Workflow Functions + +def save_workflow_chat_to_db(chat_history, workflow_name, conversation_id=None): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + + if conversation_id is None: + # Create a new conversation + conversation_name = f"{workflow_name}_Workflow_{datetime.now().strftime('%Y%m%d_%H%M%S')}" + cursor.execute(''' + INSERT INTO ChatConversations (media_id, media_name, conversation_name, created_at, updated_at) + VALUES (NULL, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP) + ''', (workflow_name, conversation_name)) + conversation_id = cursor.lastrowid + else: + # Update existing conversation + cursor.execute(''' + UPDATE ChatConversations + SET updated_at = CURRENT_TIMESTAMP + WHERE id = ? + ''', (conversation_id,)) + + # Save messages + for user_msg, ai_msg in chat_history: + if user_msg: + cursor.execute(''' + INSERT INTO ChatMessages (conversation_id, sender, message, timestamp) + VALUES (?, 'user', ?, CURRENT_TIMESTAMP) + ''', (conversation_id, user_msg)) + if ai_msg: + cursor.execute(''' + INSERT INTO ChatMessages (conversation_id, sender, message, timestamp) + VALUES (?, 'ai', ?, CURRENT_TIMESTAMP) + ''', (conversation_id, ai_msg)) + + conn.commit() + + return conversation_id, f"Chat saved successfully! Conversation ID: {conversation_id}" + except Exception as e: + logging.error(f"Error saving workflow chat to database: {str(e)}") + return None, f"Error saving chat to database: {str(e)}" + + +def get_workflow_chat(conversation_id): + """ + Retrieve a workflow chat from the database. + + Args: + conversation_id: ID of the conversation to retrieve + + Returns: + tuple: (chat_history, workflow_name, status_message) + """ + try: + with db.get_connection() as conn: + cursor = conn.cursor() + + # Get conversation details + cursor.execute(''' + SELECT media_name, conversation_name FROM ChatConversations + WHERE id = ? + ''', (conversation_id,)) + result = cursor.fetchone() + if not result: + return None, None, "Conversation not found" + + workflow_name, conversation_name = result + + # Get chat messages + cursor.execute(''' + SELECT sender, message FROM ChatMessages + WHERE conversation_id = ? + ORDER BY timestamp + ''', (conversation_id,)) + messages = cursor.fetchall() + + chat_history = [] + for sender, message in messages: + if sender == 'user': + chat_history.append((message, None)) + else: + if chat_history and chat_history[-1][1] is None: + chat_history[-1] = (chat_history[-1][0], message) + else: + chat_history.append((None, message)) + + return chat_history, workflow_name, f"Chat retrieved successfully" + except Exception as e: + logging.error(f"Error retrieving workflow chat from database: {str(e)}") + return None, None, f"Error retrieving chat from database: {str(e)}" + +# +# End of Workflow Functions +####################################################################################################################### diff --git a/App_Function_Libraries/DB/__init__.py b/App_Function_Libraries/DB/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/App_Function_Libraries/Gradio_Related.py b/App_Function_Libraries/Gradio_Related.py new file mode 100644 index 0000000000000000000000000000000000000000..e3b205f3088397d706ee40e3ba876e40d2d1bad4 --- /dev/null +++ b/App_Function_Libraries/Gradio_Related.py @@ -0,0 +1,420 @@ +# Gradio_Related.py +######################################### +# Gradio UI Functions Library +# I fucking hate Gradio. +# +######################################### +# +# Built-In Imports +import logging +import os +import webbrowser + +# +# Import 3rd-Party Libraries +import gradio as gr +# +# Local Imports +from App_Function_Libraries.DB.DB_Manager import get_db_config +from App_Function_Libraries.Gradio_UI.Arxiv_tab import create_arxiv_tab +from App_Function_Libraries.Gradio_UI.Audio_ingestion_tab import create_audio_processing_tab +from App_Function_Libraries.Gradio_UI.Book_Ingestion_tab import create_import_book_tab +from App_Function_Libraries.Gradio_UI.Character_Chat_tab import create_character_card_interaction_tab, create_character_chat_mgmt_tab, create_custom_character_card_tab, \ + create_character_card_validation_tab, create_export_characters_tab +from App_Function_Libraries.Gradio_UI.Character_interaction_tab import create_narrator_controlled_conversation_tab, \ + create_multiple_character_chat_tab +from App_Function_Libraries.Gradio_UI.Chat_ui import create_chat_management_tab, \ + create_chat_interface_four, create_chat_interface_multi_api, create_chat_interface_stacked, create_chat_interface +from App_Function_Libraries.Gradio_UI.Config_tab import create_config_editor_tab +from App_Function_Libraries.Gradio_UI.Explain_summarize_tab import create_summarize_explain_tab +from App_Function_Libraries.Gradio_UI.Export_Functionality import create_export_tab +from App_Function_Libraries.Gradio_UI.Backup_Functionality import create_backup_tab, create_view_backups_tab, \ + create_restore_backup_tab +from App_Function_Libraries.Gradio_UI.Import_Functionality import create_import_single_prompt_tab, \ + create_import_obsidian_vault_tab, create_import_item_tab, create_import_multiple_prompts_tab +from App_Function_Libraries.Gradio_UI.Introduction_tab import create_introduction_tab +from App_Function_Libraries.Gradio_UI.Keywords import create_view_keywords_tab, create_add_keyword_tab, \ + create_delete_keyword_tab, create_export_keywords_tab +from App_Function_Libraries.Gradio_UI.Live_Recording import create_live_recording_tab +from App_Function_Libraries.Gradio_UI.Llamafile_tab import create_chat_with_llamafile_tab +#from App_Function_Libraries.Gradio_UI.MMLU_Pro_tab import create_mmlu_pro_tab +from App_Function_Libraries.Gradio_UI.Media_edit import create_prompt_clone_tab, create_prompt_edit_tab, \ + create_media_edit_and_clone_tab, create_media_edit_tab +from App_Function_Libraries.Gradio_UI.Media_wiki_tab import create_mediawiki_import_tab, create_mediawiki_config_tab +from App_Function_Libraries.Gradio_UI.PDF_ingestion_tab import create_pdf_ingestion_tab, create_pdf_ingestion_test_tab +from App_Function_Libraries.Gradio_UI.Plaintext_tab_import import create_plain_text_import_tab +from App_Function_Libraries.Gradio_UI.Podcast_tab import create_podcast_tab +from App_Function_Libraries.Gradio_UI.Prompt_Suggestion_tab import create_prompt_suggestion_tab +from App_Function_Libraries.Gradio_UI.RAG_QA_Chat_tab import create_rag_qa_chat_tab, create_rag_qa_notes_management_tab, \ + create_rag_qa_chat_management_tab +from App_Function_Libraries.Gradio_UI.Re_summarize_tab import create_resummary_tab +from App_Function_Libraries.Gradio_UI.Search_Tab import create_prompt_search_tab, \ + create_search_summaries_tab, create_search_tab +from App_Function_Libraries.Gradio_UI.RAG_Chat_tab import create_rag_tab +from App_Function_Libraries.Gradio_UI.Embeddings_tab import create_embeddings_tab, create_view_embeddings_tab, \ + create_purge_embeddings_tab +from App_Function_Libraries.Gradio_UI.Trash import create_view_trash_tab, create_empty_trash_tab, \ + create_delete_trash_tab, create_search_and_mark_trash_tab +from App_Function_Libraries.Gradio_UI.Utilities import create_utilities_yt_timestamp_tab, create_utilities_yt_audio_tab, \ + create_utilities_yt_video_tab +from App_Function_Libraries.Gradio_UI.Video_transcription_tab import create_video_transcription_tab +from App_Function_Libraries.Gradio_UI.View_tab import create_manage_items_tab +from App_Function_Libraries.Gradio_UI.Website_scraping_tab import create_website_scraping_tab +from App_Function_Libraries.Gradio_UI.Chat_Workflows import chat_workflows_tab +from App_Function_Libraries.Gradio_UI.View_DB_Items_tab import create_prompt_view_tab, \ + create_view_all_with_versions_tab, create_viewing_tab +# +# Gradio UI Imports +from App_Function_Libraries.Gradio_UI.Evaluations_Benchmarks_tab import create_geval_tab, create_infinite_bench_tab +#from App_Function_Libraries.Local_LLM.Local_LLM_huggingface import create_huggingface_tab +from App_Function_Libraries.Local_LLM.Local_LLM_ollama import create_ollama_tab +# +####################################################################################################################### +# Function Definitions +# + + +# Disable Gradio Analytics +os.environ['GRADIO_ANALYTICS_ENABLED'] = 'False' + + +custom_prompt_input = None +server_mode = False +share_public = False +custom_prompt_summarize_bulleted_notes = (""" + You are a bulleted notes specialist. [INST]```When creating comprehensive bulleted notes, you should follow these guidelines: Use multiple headings based on the referenced topics, not categories like quotes or terms. Headings should be surrounded by bold formatting and not be listed as bullet points themselves. Leave no space between headings and their corresponding list items underneath. Important terms within the content should be emphasized by setting them in bold font. Any text that ends with a colon should also be bolded. Before submitting your response, review the instructions, and make any corrections necessary to adhered to the specified format. Do not reference these instructions within the notes.``` \nBased on the content between backticks create comprehensive bulleted notes.[/INST] + **Bulleted Note Creation Guidelines** + + **Headings**: + - Based on referenced topics, not categories like quotes or terms + - Surrounded by **bold** formatting + - Not listed as bullet points + - No space between headings and list items underneath + + **Emphasis**: + - **Important terms** set in bold font + - **Text ending in a colon**: also bolded + + **Review**: + - Ensure adherence to specified format + - Do not reference these instructions in your response.[INST] {{ .Prompt }} [/INST] + """) +# +# End of globals +####################################################################################################################### +# +# Start of Video/Audio Transcription and Summarization Functions +# +# Functions: +# FIXME +# +# +################################################################################################################ +# Functions for Re-Summarization +# +# Functions: +# FIXME +# End of Re-Summarization Functions +# +############################################################################################################################################################################################################################ +# +# Explain/Summarize This Tab +# +# Functions: +# FIXME +# +# +############################################################################################################################################################################################################################ +# +# Transcript Comparison Tab +# +# Functions: +# FIXME +# +# +########################################################################################################################################################################################################################### +# +# Search Tab +# +# Functions: +# FIXME +# +# End of Search Tab Functions +# +############################################################################################################################################################################################################################## +# +# Llamafile Tab +# +# Functions: +# FIXME +# +# End of Llamafile Tab Functions +############################################################################################################################################################################################################################## +# +# Chat Interface Tab Functions +# +# Functions: +# FIXME +# +# +# End of Chat Interface Tab Functions +################################################################################################################################################################################################################################ +# +# Media Edit Tab Functions +# Functions: +# Fixme +# create_media_edit_tab(): +##### Trash Tab +# FIXME +# Functions: +# +# End of Media Edit Tab Functions +################################################################################################################ +# +# Import Items Tab Functions +# +# Functions: +#FIXME +# End of Import Items Tab Functions +################################################################################################################ +# +# Export Items Tab Functions +# +# Functions: +# FIXME +# +# +# End of Export Items Tab Functions +################################################################################################################ +# +# Keyword Management Tab Functions +# +# Functions: +# create_view_keywords_tab(): +# FIXME +# +# End of Keyword Management Tab Functions +################################################################################################################ +# +# Document Editing Tab Functions +# +# Functions: +# #FIXME +# +# +################################################################################################################ +# +# Utilities Tab Functions +# Functions: +# create_utilities_yt_video_tab(): +# #FIXME + +# +# End of Utilities Tab Functions +################################################################################################################ + +# FIXME - Prompt sample box +# +# # Sample data +# prompts_category_1 = [ +# "What are the key points discussed in the video?", +# "Summarize the main arguments made by the speaker.", +# "Describe the conclusions of the study presented." +# ] +# +# prompts_category_2 = [ +# "How does the proposed solution address the problem?", +# "What are the implications of the findings?", +# "Can you explain the theory behind the observed phenomenon?" +# ] +# +# all_prompts2 = prompts_category_1 + prompts_category_2 + + +def launch_ui(share_public=None, server_mode=False): + webbrowser.open_new_tab('http://127.0.0.1:7860/?__theme=dark') + share=share_public + css = """ + .result-box { + margin-bottom: 20px; + border: 1px solid #ddd; + padding: 10px; + } + .result-box.error { + border-color: #ff0000; + background-color: #ffeeee; + } + .transcription, .summary { + max-height: 800px; + overflow-y: auto; + border: 1px solid #eee; + padding: 10px; + margin-top: 10px; + } + """ + + with gr.Blocks(theme='bethecloud/storj_theme',css=css) as iface: + gr.HTML( + """ + + """ + ) + db_config = get_db_config() + db_type = db_config['type'] + gr.Markdown(f"# tl/dw: Your LLM-powered Research Multi-tool") + gr.Markdown(f"(Using {db_type.capitalize()} Database)") + with gr.Tabs(): + with gr.TabItem("Transcription / Summarization / Ingestion", id="ingestion-grouping", visible=True): + with gr.Tabs(): + create_video_transcription_tab() + create_audio_processing_tab() + create_podcast_tab() + create_import_book_tab() + create_plain_text_import_tab() + create_website_scraping_tab() + create_pdf_ingestion_tab() + create_pdf_ingestion_test_tab() + create_resummary_tab() + create_summarize_explain_tab() + create_live_recording_tab() + create_arxiv_tab() + + with gr.TabItem("Text Search", id="text search", visible=True): + create_search_tab() + create_search_summaries_tab() + + with gr.TabItem("RAG Chat/Search", id="RAG Chat Notes group", visible=True): + create_rag_tab() + create_rag_qa_chat_tab() + create_rag_qa_notes_management_tab() + create_rag_qa_chat_management_tab() + + with gr.TabItem("Chat with an LLM", id="LLM Chat group", visible=True): + create_chat_interface() + create_chat_interface_stacked() + create_chat_interface_multi_api() + create_chat_interface_four() + create_chat_with_llamafile_tab() + create_chat_management_tab() + chat_workflows_tab() + + + with gr.TabItem("Character Chat", id="character chat group", visible=True): + create_character_card_interaction_tab() + create_character_chat_mgmt_tab() + create_custom_character_card_tab() + create_character_card_validation_tab() + create_multiple_character_chat_tab() + create_narrator_controlled_conversation_tab() + create_export_characters_tab() + + + with gr.TabItem("View DB Items", id="view db items group", visible=True): + # This one works + create_view_all_with_versions_tab() + # This one is WIP + create_viewing_tab() + create_prompt_view_tab() + + + with gr.TabItem("Prompts", id='view prompts group', visible=True): + create_prompt_view_tab() + create_prompt_search_tab() + create_prompt_edit_tab() + create_prompt_clone_tab() + create_prompt_suggestion_tab() + + + with gr.TabItem("Manage / Edit Existing Items", id="manage group", visible=True): + create_media_edit_tab() + create_manage_items_tab() + create_media_edit_and_clone_tab() + # FIXME + #create_compare_transcripts_tab() + + + with gr.TabItem("Embeddings Management", id="embeddings group", visible=True): + create_embeddings_tab() + create_view_embeddings_tab() + create_purge_embeddings_tab() + + with gr.TabItem("Writing Tools", id="writing_tools group", visible=True): + from App_Function_Libraries.Gradio_UI.Writing_tab import create_document_feedback_tab + create_document_feedback_tab() + from App_Function_Libraries.Gradio_UI.Writing_tab import create_grammar_style_check_tab + create_grammar_style_check_tab() + from App_Function_Libraries.Gradio_UI.Writing_tab import create_tone_adjustment_tab + create_tone_adjustment_tab() + from App_Function_Libraries.Gradio_UI.Writing_tab import create_creative_writing_tab + create_creative_writing_tab() + from App_Function_Libraries.Gradio_UI.Writing_tab import create_mikupad_tab + create_mikupad_tab() + + + with gr.TabItem("Keywords", id="keywords group", visible=True): + create_view_keywords_tab() + create_add_keyword_tab() + create_delete_keyword_tab() + create_export_keywords_tab() + + with gr.TabItem("Import", id="import group", visible=True): + create_import_item_tab() + create_import_obsidian_vault_tab() + create_import_single_prompt_tab() + create_import_multiple_prompts_tab() + create_mediawiki_import_tab() + create_mediawiki_config_tab() + + with gr.TabItem("Export", id="export group", visible=True): + create_export_tab() + + with gr.TabItem("Backup Management", id="backup group", visible=True): + create_backup_tab() + create_view_backups_tab() + create_restore_backup_tab() + + with gr.TabItem("Utilities", id="util group", visible=True): + create_utilities_yt_video_tab() + create_utilities_yt_audio_tab() + create_utilities_yt_timestamp_tab() + + with gr.TabItem("Local LLM", id="local llm group", visible=True): + create_chat_with_llamafile_tab() + create_ollama_tab() + #create_huggingface_tab() + + with gr.TabItem("Trashcan", id="trashcan group", visible=True): + create_search_and_mark_trash_tab() + create_view_trash_tab() + create_delete_trash_tab() + create_empty_trash_tab() + + with gr.TabItem("Evaluations", id="eval", visible=True): + create_geval_tab() + create_infinite_bench_tab() + # FIXME + #create_mmlu_pro_tab() + + with gr.TabItem("Introduction/Help", id="introduction group", visible=True): + create_introduction_tab() + + with gr.TabItem("Config Editor", id="config group"): + create_config_editor_tab() + + # Launch the interface + server_port_variable = 7860 + os.environ['GRADIO_ANALYTICS_ENABLED'] = 'False' + if share==True: + iface.launch(share=True) + elif server_mode and not share_public: + iface.launch(share=False, server_name="0.0.0.0", server_port=server_port_variable, ) + else: + try: + iface.launch(share=False, server_name="0.0.0.0", server_port=server_port_variable, ) + except Exception as e: + logging.error(f"Error launching interface: {str(e)}") diff --git a/App_Function_Libraries/Gradio_UI/Arxiv_tab.py b/App_Function_Libraries/Gradio_UI/Arxiv_tab.py new file mode 100644 index 0000000000000000000000000000000000000000..01c1222f73e8104d2ac6de25e6450900ed9d96ec --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Arxiv_tab.py @@ -0,0 +1,230 @@ +# Arxiv_tab.py +# Description: This file contains the Gradio UI for searching, browsing, and ingesting arXiv papers. +# +# Imports +import tempfile +from datetime import datetime +import requests + +from App_Function_Libraries.PDF.PDF_Ingestion_Lib import extract_text_and_format_from_pdf +# +# Local Imports +from App_Function_Libraries.Third_Party.Arxiv import convert_xml_to_markdown, fetch_arxiv_xml, parse_arxiv_feed, \ + build_query_url, ARXIV_PAGE_SIZE, fetch_arxiv_pdf_url +from App_Function_Libraries.DB.DB_Manager import add_media_with_keywords +# +import gradio as gr +# +##################################################################################################### +# +# Functions: + +def create_arxiv_tab(): + with gr.TabItem("Arxiv Search & Ingest", visible=True): + gr.Markdown("# arXiv Search, Browse, Download, and Ingest") + gr.Markdown("#### Thank you to arXiv for use of its open access interoperability.") + with gr.Row(): + with gr.Column(scale=1): + # Search Inputs + with gr.Row(): + with gr.Column(): + search_query = gr.Textbox(label="Search Query", placeholder="e.g., machine learning") + author_filter = gr.Textbox(label="Author", placeholder="e.g., John Doe") + year_filter = gr.Number(label="Year", precision=0) + search_button = gr.Button("Search") + + with gr.Column(scale=2): + # Pagination Controls + paper_selector = gr.Radio(label="Select a Paper", choices=[], interactive=True) + prev_button = gr.Button("Previous Page") + next_button = gr.Button("Next Page") + page_info = gr.Textbox(label="Page", value="1", interactive=False) + + # Ingestion Section + with gr.Row(): + with gr.Column(): + # Paper Details View + paper_view = gr.Markdown(label="Paper Details") + arxiv_keywords = gr.Textbox(label="Additional Keywords (comma-separated)", + placeholder="e.g., AI, Deep Learning") + ingest_button = gr.Button("Ingest Selected Paper") + ingest_result = gr.Textbox(label="Ingestion Result", interactive=False) + + # Define States for Pagination and Selection + state = gr.State(value={"start": 0, "current_page": 1, "last_query": None, "entries": []}) + selected_paper_id = gr.State(value=None) + + def search_arxiv(query, author, year): + start = 0 + url = build_query_url(query, author, year, start) + try: + response = requests.get(url) + response.raise_for_status() + except requests.exceptions.RequestException as e: + return gr.update(value=[]), gr.update(value=f"**Error:** {str(e)}"), state.value + + entries = parse_arxiv_feed(response.text) + state.value = {"start": start, "current_page": 1, "last_query": (query, author, year), "entries": entries} + if not entries: + return gr.update(value=[]), "No results found.", state.value + + # Update the dropdown with paper titles for selection + titles = [entry['title'] for entry in entries] + return gr.update(choices=titles), "1", state.value + + # Dead code? FIXME + def handle_pagination(direction): + current_state = state.value + query, author, year = current_state["last_query"] + new_page = current_state["current_page"] + direction + if new_page < 1: + new_page = 1 + start = (new_page - 1) * ARXIV_PAGE_SIZE + url = build_query_url(query, author, year, start) + try: + response = requests.get(url) + response.raise_for_status() + except requests.exceptions.RequestException as e: + return gr.update(), gr.update() + + entries = parse_arxiv_feed(response.text) + if entries: + current_state["start"] = start + current_state["current_page"] = new_page + current_state["entries"] = entries + state.value = current_state + + # Update the dropdown with paper titles for the new page + titles = [entry['title'] for entry in entries] + return gr.update(choices=titles), str(new_page) + else: + # If no entries, do not change the page + return gr.update(), gr.update() + + def load_selected_paper(selected_title): + if not selected_title: + return "Please select a paper to view." + + # Find the selected paper from state + for entry in state.value["entries"]: + if entry['title'] == selected_title: + paper_id = entry['id'] + break + else: + return "Paper not found." + + try: + # Fetch the PDF URL and download the full-text + pdf_url = fetch_arxiv_pdf_url(paper_id) + response = requests.get(pdf_url) + response.raise_for_status() + + # Save the PDF temporarily + with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_pdf: + temp_pdf.write(response.content) + temp_pdf_path = temp_pdf.name + + # Convert PDF to markdown using your PDF ingestion function + full_text_markdown = extract_text_and_format_from_pdf(temp_pdf_path) + + selected_paper_id.value = paper_id + return full_text_markdown + except Exception as e: + return f"Error loading full paper: {str(e)}" + + def process_and_ingest_arxiv_paper(paper_id, additional_keywords): + try: + if not paper_id: + return "Please select a paper to ingest." + + # Fetch the PDF URL + pdf_url = fetch_arxiv_pdf_url(paper_id) + + # Download the PDF + response = requests.get(pdf_url) + response.raise_for_status() + + # Save the PDF temporarily + with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_pdf: + temp_pdf.write(response.content) + temp_pdf_path = temp_pdf.name + + # Convert PDF to markdown using your PDF ingestion function + markdown_text = extract_text_and_format_from_pdf(temp_pdf_path) + + # Fetch metadata from arXiv to get title, authors, and categories + xml_content = fetch_arxiv_xml(paper_id) + _, title, authors, categories = convert_xml_to_markdown(xml_content) + + # Prepare the arXiv paper URL for access/download + paper_url = f"https://arxiv.org/abs/{paper_id}" + + # Prepare the keywords for ingestion + keywords = f"arxiv,{','.join(categories)}" + if additional_keywords: + keywords += f",{additional_keywords}" + + # Ingest full paper markdown content + add_media_with_keywords( + url=paper_url, + title=title, + media_type='document', + content=markdown_text, # Full paper content in markdown + keywords=keywords, + prompt='No prompt for arXiv papers', + summary='Full arXiv paper ingested from PDF', + transcription_model='None', + author=', '.join(authors), + ingestion_date=datetime.now().strftime('%Y-%m-%d') + ) + + # Return success message with paper title and authors + return f"arXiv paper '{title}' by {', '.join(authors)} ingested successfully." + except Exception as e: + # Return error message if anything goes wrong + return f"Error processing arXiv paper: {str(e)}" + + # Event Handlers + # Connect Search Button + search_button.click( + fn=search_arxiv, + inputs=[search_query, author_filter, year_filter], + outputs=[paper_selector, page_info, state], + queue=True + ) + + # Connect Next Button + next_button.click( + fn=lambda: handle_pagination(1), + inputs=None, + outputs=[paper_selector, page_info], + queue=True + ) + + # Connect Previous Button + prev_button.click( + fn=lambda: handle_pagination(-1), + inputs=None, + outputs=[paper_selector, page_info], + queue=True + ) + + # When the user selects a paper in the Dropdown + paper_selector.change( + fn=load_selected_paper, + inputs=paper_selector, + outputs=paper_view, + queue=True + ) + + # Connect Ingest Button + ingest_button.click( + fn=process_and_ingest_arxiv_paper, + inputs=[selected_paper_id, arxiv_keywords], + outputs=ingest_result, + queue=True + ) + +# +# End of File +##################################################################################################### diff --git a/App_Function_Libraries/Gradio_UI/Audio_ingestion_tab.py b/App_Function_Libraries/Gradio_UI/Audio_ingestion_tab.py new file mode 100644 index 0000000000000000000000000000000000000000..3eee842b05c12075fcb23f9ee7b623c6f768604a --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Audio_ingestion_tab.py @@ -0,0 +1,167 @@ +# Audio_ingestion_tab.py +# Description: Gradio UI for ingesting audio files into the database +# +# Imports +# +# External Imports +import gradio as gr +# +# Local Imports +from App_Function_Libraries.Audio.Audio_Files import process_audio_files +from App_Function_Libraries.DB.DB_Manager import load_preset_prompts +from App_Function_Libraries.Gradio_UI.Chat_ui import update_user_prompt +from App_Function_Libraries.Gradio_UI.Gradio_Shared import whisper_models +from App_Function_Libraries.Utils.Utils import cleanup_temp_files +# Import metrics logging +from App_Function_Libraries.Metrics.metrics_logger import log_counter, log_histogram +from App_Function_Libraries.Metrics.logger_config import logger +# +####################################################################################################################### +# Functions: + +def create_audio_processing_tab(): + with gr.TabItem("Audio File Transcription + Summarization", visible=True): + gr.Markdown("# Transcribe & Summarize Audio Files from URLs or Local Files!") + with gr.Row(): + with gr.Column(): + audio_url_input = gr.Textbox(label="Audio File URL(s)", placeholder="Enter the URL(s) of the audio file(s), one per line") + audio_file_input = gr.File(label="Upload Audio File", file_types=["audio/*"]) + custom_title_input = gr.Textbox(label="Custom Title/Name", placeholder="Enter a custom title or name for the audio file") + use_cookies_input = gr.Checkbox(label="Use cookies for authenticated download", value=False) + cookies_input = gr.Textbox( + label="Audio Download Cookies", + placeholder="Paste your cookies here (JSON format)", + lines=3, + visible=False + ) + + use_cookies_input.change( + fn=lambda x: gr.update(visible=x), + inputs=[use_cookies_input], + outputs=[cookies_input] + ) + + diarize_input = gr.Checkbox(label="Enable Speaker Diarization", value=False) + whisper_model_input = gr.Dropdown(choices=whisper_models, value="medium", label="Whisper Model") + keep_timestamps_input = gr.Checkbox(label="Keep Timestamps", value=True) + + with gr.Row(): + custom_prompt_checkbox = gr.Checkbox(label="Use a Custom Prompt", + value=False, + visible=True) + preset_prompt_checkbox = gr.Checkbox(label="Use a pre-set Prompt", + value=False, + visible=True) + with gr.Row(): + preset_prompt = gr.Dropdown(label="Select Preset Prompt", + choices=load_preset_prompts(), + visible=False) + with gr.Row(): + custom_prompt_input = gr.Textbox(label="Custom Prompt", + placeholder="Enter custom prompt here", + lines=3, + visible=False) + with gr.Row(): + system_prompt_input = gr.Textbox(label="System Prompt", + value="""You are a bulleted notes specialist. [INST]```When creating comprehensive bulleted notes, you should follow these guidelines: Use multiple headings based on the referenced topics, not categories like quotes or terms. Headings should be surrounded by bold formatting and not be listed as bullet points themselves. Leave no space between headings and their corresponding list items underneath. Important terms within the content should be emphasized by setting them in bold font. Any text that ends with a colon should also be bolded. Before submitting your response, review the instructions, and make any corrections necessary to adhered to the specified format. Do not reference these instructions within the notes.``` \nBased on the content between backticks create comprehensive bulleted notes.[/INST] +**Bulleted Note Creation Guidelines** + +**Headings**: +- Based on referenced topics, not categories like quotes or terms +- Surrounded by **bold** formatting +- Not listed as bullet points +- No space between headings and list items underneath + +**Emphasis**: +- **Important terms** set in bold font +- **Text ending in a colon**: also bolded + +**Review**: +- Ensure adherence to specified format +- Do not reference these instructions in your response.[INST] {{ .Prompt }} [/INST] +""", + lines=3, + visible=False) + + custom_prompt_checkbox.change( + fn=lambda x: (gr.update(visible=x), gr.update(visible=x)), + inputs=[custom_prompt_checkbox], + outputs=[custom_prompt_input, system_prompt_input] + ) + preset_prompt_checkbox.change( + fn=lambda x: gr.update(visible=x), + inputs=[preset_prompt_checkbox], + outputs=[preset_prompt] + ) + + def update_prompts(preset_name): + prompts = update_user_prompt(preset_name) + return ( + gr.update(value=prompts["user_prompt"], visible=True), + gr.update(value=prompts["system_prompt"], visible=True) + ) + + preset_prompt.change( + update_prompts, + inputs=preset_prompt, + outputs=[custom_prompt_input, system_prompt_input] + ) + + api_name_input = gr.Dropdown( + choices=[None, "Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", "DeepSeek", "Mistral", "OpenRouter", + "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM","ollama", "HuggingFace", "Custom-OpenAI-API"], + value=None, + label="API for Summarization (Optional)" + ) + api_key_input = gr.Textbox(label="API Key (if required)", placeholder="Enter your API key here", type="password") + custom_keywords_input = gr.Textbox(label="Custom Keywords", placeholder="Enter custom keywords, comma-separated") + keep_original_input = gr.Checkbox(label="Keep original audio file", value=False) + + chunking_options_checkbox = gr.Checkbox(label="Show Chunking Options", value=False) + with gr.Row(visible=False) as chunking_options_box: + gr.Markdown("### Chunking Options") + with gr.Column(): + chunk_method = gr.Dropdown(choices=['words', 'sentences', 'paragraphs', 'tokens'], label="Chunking Method") + max_chunk_size = gr.Slider(minimum=100, maximum=1000, value=300, step=50, label="Max Chunk Size") + chunk_overlap = gr.Slider(minimum=0, maximum=100, value=0, step=10, label="Chunk Overlap") + use_adaptive_chunking = gr.Checkbox(label="Use Adaptive Chunking") + use_multi_level_chunking = gr.Checkbox(label="Use Multi-level Chunking") + chunk_language = gr.Dropdown(choices=['english', 'french', 'german', 'spanish'], label="Chunking Language") + + chunking_options_checkbox.change( + fn=lambda x: gr.update(visible=x), + inputs=[chunking_options_checkbox], + outputs=[chunking_options_box] + ) + + process_audio_button = gr.Button("Process Audio File(s)") + + with gr.Column(): + audio_progress_output = gr.Textbox(label="Progress") + audio_transcription_output = gr.Textbox(label="Transcription") + audio_summary_output = gr.Textbox(label="Summary") + download_transcription = gr.File(label="Download All Transcriptions as JSON") + download_summary = gr.File(label="Download All Summaries as Text") + + process_audio_button.click( + fn=process_audio_files, + inputs=[audio_url_input, audio_file_input, whisper_model_input, api_name_input, api_key_input, + use_cookies_input, cookies_input, keep_original_input, custom_keywords_input, custom_prompt_input, + chunk_method, max_chunk_size, chunk_overlap, use_adaptive_chunking, use_multi_level_chunking, + chunk_language, diarize_input, keep_timestamps_input, custom_title_input], + outputs=[audio_progress_output, audio_transcription_output, audio_summary_output] + ) + + def on_file_clear(file): + if file is None: + cleanup_temp_files() + + audio_file_input.clear( + fn=on_file_clear, + inputs=[audio_file_input], + outputs=[] + ) + +# +# End of Audio_ingestion_tab.py +####################################################################################################################### \ No newline at end of file diff --git a/App_Function_Libraries/Gradio_UI/Backup_Functionality.py b/App_Function_Libraries/Gradio_UI/Backup_Functionality.py new file mode 100644 index 0000000000000000000000000000000000000000..c4bc198ec7ea0b811e7d60bf0756f305bc4d3951 --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Backup_Functionality.py @@ -0,0 +1,71 @@ +# Backup_Functionality.py +# Functionality for exporting items as markdown files +# +# Imports: +import os +import shutil +import gradio as gr +# +# Local Imports: +from App_Function_Libraries.DB.DB_Manager import create_automated_backup, db_path, backup_dir +# +# End of Imports +####################################################################################################################### +# +# Functions: + +def create_backup(): + backup_file = create_automated_backup(db_path, backup_dir) + return f"Backup created: {backup_file}" + + +def list_backups(): + backups = [f for f in os.listdir(backup_dir) if f.endswith('.db')] + return "\n".join(backups) + + +def restore_backup(backup_name: str) -> str: + backup_path_location: str = os.path.join(str(backup_dir), backup_name) + if os.path.exists(backup_path_location): + shutil.copy2(str(backup_path_location), str(db_path)) + return f"Database restored from {backup_name}" + else: + return "Backup file not found" + + +def create_backup_tab(): + with gr.Tab("Create Backup", visible=True): + gr.Markdown("# Create a backup of the database") + gr.Markdown("This will create a backup of the database in the backup directory(the default backup directory is `/tldw_DB_Backups/')") + with gr.Row(): + with gr.Column(): + create_button = gr.Button("Create Backup") + create_output = gr.Textbox(label="Result") + with gr.Column(): + create_button.click(create_backup, inputs=[], outputs=create_output) + + +def create_view_backups_tab(): + with gr.TabItem("View Backups", visible=True): + gr.Markdown("# Browse available backups") + with gr.Row(): + with gr.Column(): + view_button = gr.Button("View Backups") + with gr.Column(): + backup_list = gr.Textbox(label="Available Backups") + view_button.click(list_backups, inputs=[], outputs=backup_list) + + +def create_restore_backup_tab(): + with gr.TabItem("Restore Backup", visible=True): + gr.Markdown("# Restore a backup of the database") + with gr.Column(): + backup_input = gr.Textbox(label="Backup Filename") + restore_button = gr.Button("Restore") + with gr.Column(): + restore_output = gr.Textbox(label="Result") + restore_button.click(restore_backup, inputs=[backup_input], outputs=restore_output) + +# +# End of Functions +####################################################################################################################### diff --git a/App_Function_Libraries/Gradio_UI/Book_Ingestion_tab.py b/App_Function_Libraries/Gradio_UI/Book_Ingestion_tab.py new file mode 100644 index 0000000000000000000000000000000000000000..cc455dfa67109a7f9ab95ab17fe1d67ae9142b67 --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Book_Ingestion_tab.py @@ -0,0 +1,100 @@ +# Book_Ingestion_tab.py +# Functionality to import epubs/ebooks into the system. +#################### +# Function List +# +# 1. create_import_book_tab() +# 2. import_epub(epub_file, title, author, keywords, system_prompt, user_prompt, auto_summarize, api_name, api_key) +# +#################### +# Imports +# +# External Imports +import gradio as gr +# +# Local Imports +from App_Function_Libraries.Books.Book_Ingestion_Lib import process_zip_file, import_epub, import_file_handler +# +######################################################################################################################## +# +# Functions: + + + +def create_import_book_tab(): + with gr.TabItem("Ebook(epub) Files", visible=True): + with gr.Row(): + with gr.Column(): + gr.Markdown("# Import .epub files") + gr.Markdown("Upload a single .epub file or a .zip file containing multiple .epub files") + gr.Markdown( + "🔗 **How to remove DRM from your ebooks:** [Reddit Guide](https://www.reddit.com/r/Calibre/comments/1ck4w8e/2024_guide_on_removing_drm_from_kobo_kindle_ebooks/)") + import_file = gr.File(label="Upload file for import", file_types=[".epub", ".zip"]) + title_input = gr.Textbox(label="Title", placeholder="Enter the title of the content (for single files)") + author_input = gr.Textbox(label="Author", placeholder="Enter the author's name (for single files)") + keywords_input = gr.Textbox(label="Keywords (like genre or publish year)", + placeholder="Enter keywords, comma-separated") + system_prompt_input = gr.Textbox(label="System Prompt", lines=3, + value="""" + You are a bulleted notes specialist. [INST]```When creating comprehensive bulleted notes, you should follow these guidelines: Use multiple headings based on the referenced topics, not categories like quotes or terms. Headings should be surrounded by bold formatting and not be listed as bullet points themselves. Leave no space between headings and their corresponding list items underneath. Important terms within the content should be emphasized by setting them in bold font. Any text that ends with a colon should also be bolded. Before submitting your response, review the instructions, and make any corrections necessary to adhered to the specified format. Do not reference these instructions within the notes.``` \nBased on the content between backticks create comprehensive bulleted notes.[/INST] + **Bulleted Note Creation Guidelines** + + **Headings**: + - Based on referenced topics, not categories like quotes or terms + - Surrounded by **bold** formatting + - Not listed as bullet points + - No space between headings and list items underneath + + **Emphasis**: + - **Important terms** set in bold font + - **Text ending in a colon**: also bolded + + **Review**: + - Ensure adherence to specified format + - Do not reference these instructions in your response.[INST] {{ .Prompt }} [/INST] + """, ) + custom_prompt_input = gr.Textbox(label="Custom User Prompt", + placeholder="Enter a custom user prompt for summarization (optional)") + auto_summarize_checkbox = gr.Checkbox(label="Auto-summarize", value=False) + api_name_input = gr.Dropdown( + choices=[None, "Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", "DeepSeek", "Mistral", + "OpenRouter", "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM", "ollama", "HuggingFace"], + label="API for Auto-summarization" + ) + api_key_input = gr.Textbox(label="API Key", type="password") + + # Chunking options + max_chunk_size = gr.Slider(minimum=100, maximum=2000, value=500, step=50, label="Max Chunk Size") + chunk_overlap = gr.Slider(minimum=0, maximum=500, value=200, step=10, label="Chunk Overlap") + custom_chapter_pattern = gr.Textbox(label="Custom Chapter Pattern (optional)", + placeholder="Enter a custom regex pattern for chapter detection") + + + import_button = gr.Button("Import eBook(s)") + with gr.Column(): + with gr.Row(): + import_output = gr.Textbox(label="Import Status", lines=10, interactive=False) + + import_button.click( + fn=import_file_handler, + inputs=[ + import_file, + title_input, + author_input, + keywords_input, + custom_prompt_input, + auto_summarize_checkbox, + api_name_input, + api_key_input, + max_chunk_size, + chunk_overlap, + custom_chapter_pattern + ], + outputs=import_output + ) + + return import_file, title_input, author_input, keywords_input, system_prompt_input, custom_prompt_input, auto_summarize_checkbox, api_name_input, api_key_input, import_button, import_output + +# +# End of File +######################################################################################################################## \ No newline at end of file diff --git a/App_Function_Libraries/Gradio_UI/Character_Chat_tab.py b/App_Function_Libraries/Gradio_UI/Character_Chat_tab.py new file mode 100644 index 0000000000000000000000000000000000000000..86c173ea2961306aea6edb60c29d62cd2b4decf0 --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Character_Chat_tab.py @@ -0,0 +1,1846 @@ +# Character_Interaction_Library.py +# Description: Library for character card import functions +# +# Imports +import re +import tempfile +import uuid +from datetime import datetime +import json +import logging +import io +import base64 +from typing import Dict, Any, Optional, List, Tuple, Union, cast +import zipfile +# +# External Imports +from PIL import Image +import gradio as gr +# +# Local Imports +from App_Function_Libraries.Character_Chat.Character_Chat_Lib import validate_character_book, validate_v2_card, \ + replace_placeholders, replace_user_placeholder, extract_json_from_image, parse_character_book, \ + load_chat_and_character, load_chat_history, load_character_and_image, extract_character_id, load_character_wrapper +from App_Function_Libraries.Chat import chat +from App_Function_Libraries.DB.Character_Chat_DB import ( + add_character_card, + get_character_cards, + get_character_card_by_id, + add_character_chat, + get_character_chats, + get_character_chat_by_id, + update_character_chat, + delete_character_chat, + delete_character_card, + update_character_card, search_character_chats, +) +from App_Function_Libraries.Utils.Utils import sanitize_user_input +# +############################################################################################################ +# +# Functions: + +################################################################################# +# +# Character card import functions: + +def import_character_card(file): + if file is None: + return None, gr.update(), "No file provided for character card import" + + try: + if file.name.lower().endswith(('.png', '.webp')): + json_data = extract_json_from_image(file) + if not json_data: + return None, gr.update(), "No character card data found in the image. This might not be a valid character card image." + elif file.name.lower().endswith('.json'): + with open(file.name, 'r', encoding='utf-8') as f: + json_data = f.read() + else: + return None, gr.update(), "Unsupported file type. Please upload a PNG/WebP image or a JSON file." + + card_data = import_character_card_json(json_data) + if not card_data: + return None, gr.update(), "Failed to parse character card data. The file might not contain valid character information." + + # Save image data for PNG/WebP files + if file.name.lower().endswith(('.png', '.webp')): + with Image.open(file) as img: + img_byte_arr = io.BytesIO() + img.save(img_byte_arr, format='PNG') + card_data['image'] = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8') + + # Save character card to database + character_id = add_character_card(card_data) + if character_id: + characters = get_character_cards() + character_names = [char['name'] for char in characters] + return card_data, gr.update( + choices=character_names), f"Character card '{card_data['name']}' imported successfully." + else: + return None, gr.update(), f"Failed to save character card '{card_data.get('name', 'Unknown')}'. It may already exist." + except Exception as e: + logging.error(f"Error importing character card: {e}") + return None, gr.update(), f"Error importing character card: {e}" + + +def import_character_card_json(json_content: str) -> Optional[Dict[str, Any]]: + try: + json_content = json_content.strip() + card_data = json.loads(json_content) + + if 'spec' in card_data and card_data['spec'] == 'chara_card_v2': + logging.info("Detected V2 character card") + return parse_v2_card(card_data) + else: + logging.info("Assuming V1 character card") + return parse_v1_card(card_data) + except json.JSONDecodeError as e: + logging.error(f"JSON decode error: {e}") + except Exception as e: + logging.error(f"Unexpected error parsing JSON: {e}") + return None + + + +def parse_v2_card(card_data: Dict[str, Any]) -> Optional[Dict[str, Any]]: + try: + # Validate spec_version + if card_data.get('spec_version') != '2.0': + logging.warning(f"Unsupported V2 spec version: {card_data.get('spec_version')}") + return None + + data = card_data['data'] + + # Ensure all required fields are present + required_fields = ['name', 'description', 'personality', 'scenario', 'first_mes', 'mes_example'] + for field in required_fields: + if field not in data: + logging.error(f"Missing required field in V2 card: {field}") + return None + + # Handle new V2 fields + parsed_data = { + 'name': data['name'], + 'description': data['description'], + 'personality': data['personality'], + 'scenario': data['scenario'], + 'first_mes': data['first_mes'], + 'mes_example': data['mes_example'], + 'creator_notes': data.get('creator_notes', ''), + 'system_prompt': data.get('system_prompt', ''), + 'post_history_instructions': data.get('post_history_instructions', ''), + 'alternate_greetings': data.get('alternate_greetings', []), + 'tags': data.get('tags', []), + 'creator': data.get('creator', ''), + 'character_version': data.get('character_version', ''), + 'extensions': data.get('extensions', {}) + } + + # Handle character_book if present + if 'character_book' in data: + parsed_data['character_book'] = parse_character_book(data['character_book']) + + return parsed_data + except KeyError as e: + logging.error(f"Missing key in V2 card structure: {e}") + except Exception as e: + logging.error(f"Error parsing V2 card: {e}") + return None + +def parse_v1_card(card_data: Dict[str, Any]) -> Dict[str, Any]: + # Ensure all required V1 fields are present + required_fields = ['name', 'description', 'personality', 'scenario', 'first_mes', 'mes_example'] + for field in required_fields: + if field not in card_data: + logging.error(f"Missing required field in V1 card: {field}") + raise ValueError(f"Missing required field in V1 card: {field}") + + # Convert V1 to V2 format + v2_data: Dict[str, Union[str, List[str], Dict[str, Any]]] = { + 'name': card_data['name'], + 'description': card_data['description'], + 'personality': card_data['personality'], + 'scenario': card_data['scenario'], + 'first_mes': card_data['first_mes'], + 'mes_example': card_data['mes_example'], + 'creator_notes': cast(str, card_data.get('creator_notes', '')), + 'system_prompt': cast(str, card_data.get('system_prompt', '')), + 'post_history_instructions': cast(str, card_data.get('post_history_instructions', '')), + 'alternate_greetings': cast(List[str], card_data.get('alternate_greetings', [])), + 'tags': cast(List[str], card_data.get('tags', [])), + 'creator': cast(str, card_data.get('creator', '')), + 'character_version': cast(str, card_data.get('character_version', '')), + 'extensions': {} + } + + # Move any non-standard V1 fields to extensions + for key, value in card_data.items(): + if key not in v2_data: + v2_data['extensions'][key] = value + + return v2_data + +# +# End of Character card import functions +#################################################### + +#################################################### +# +# Character card export functions + +def export_character_as_json(character_id): + character = get_character_card_by_id(character_id) + if character: + # Remove the 'id' field from the character data + character_data = {k: v for k, v in character.items() if k != 'id'} + + # Convert image to base64 if it exists + if 'image' in character_data and character_data['image']: + image_data = base64.b64decode(character_data['image']) + img = Image.open(io.BytesIO(image_data)) + buffered = io.BytesIO() + img.save(buffered, format="PNG") + character_data['image'] = base64.b64encode(buffered.getvalue()).decode('utf-8') + + json_data = json.dumps(character_data, indent=2) + return json_data + return None + +def export_all_characters_as_zip(): + characters = get_character_cards() + with tempfile.NamedTemporaryFile(mode='wb', delete=False, suffix='.zip') as temp_zip: + with zipfile.ZipFile(temp_zip, 'w') as zf: + for character in characters: + character_data = {k: v for k, v in character.items() if k != 'id'} + + # Convert image to base64 if it exists + if 'image' in character_data and character_data['image']: + image_data = base64.b64decode(character_data['image']) + img = Image.open(io.BytesIO(image_data)) + buffered = io.BytesIO() + img.save(buffered, format="PNG") + character_data['image'] = base64.b64encode(buffered.getvalue()).decode('utf-8') + json_data = json.dumps(character_data, indent=2) + zf.writestr(f"{character['name']}.json", json_data) + return temp_zip.name + +def export_single_character(character_selection): + if not character_selection: + return None, "No character selected." + + character_id = int(character_selection.split('(ID: ')[1].rstrip(')')) + json_data = export_character_as_json(character_id) + + if json_data: + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', encoding='utf-8') as temp_file: + temp_file.write(json_data) + return temp_file.name, f"Character '{character_selection.split(' (ID:')[0]}' exported successfully." + else: + return None, f"Failed to export character '{character_selection.split(' (ID:')[0]}'." + +def export_all_characters(): + zip_path = export_all_characters_as_zip() + return zip_path, "All characters exported successfully." + +# +# End of Character card export functions +#################################################### + +#################################################### +# +# Gradio tabs + +def create_character_card_interaction_tab(): + with gr.TabItem("Chat with a Character Card", visible=True): + gr.Markdown("# Chat with a Character Card") + with gr.Row(): + with gr.Column(scale=1): + character_image = gr.Image(label="Character Image", type="pil") + character_card_upload = gr.File( + label="Upload Character Card (PNG, WEBP, JSON)", + file_types=[".png", ".webp", ".json"] + ) + import_card_button = gr.Button("Import Character Card") + load_characters_button = gr.Button("Load Existing Characters") + character_dropdown = gr.Dropdown(label="Select Character", choices=[]) + user_name_input = gr.Textbox(label="Your Name", placeholder="Enter your name here") + api_name_input = gr.Dropdown( + choices=[ + "Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", "DeepSeek", "Mistral", + "OpenRouter", "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM", "ollama", "HuggingFace", + "Custom-OpenAI-API" + ], + value="HuggingFace", + label="API for Interaction (Mandatory)" + ) + api_key_input = gr.Textbox( + label="API Key (if not set in Config_Files/config.txt)", + placeholder="Enter your API key here", type="password" + ) + temperature_slider = gr.Slider( + minimum=0.0, maximum=2.0, value=0.7, step=0.05, label="Temperature" + ) + import_chat_button = gr.Button("Import Chat History") + chat_file_upload = gr.File(label="Upload Chat History JSON", visible=True) + + # Chat History Import and Search + gr.Markdown("## Search and Load Existing Chats") + chat_search_query = gr.Textbox( + label="Search Chats", + placeholder="Enter chat name or keywords to search" + ) + chat_search_button = gr.Button("Search Chats") + chat_search_dropdown = gr.Dropdown(label="Search Results", choices=[], visible=False) + load_chat_button = gr.Button("Load Selected Chat", visible=False) + + # Checkbox to Decide Whether to Save Chats by Default + auto_save_checkbox = gr.Checkbox(label="Save chats automatically", value=True) + chat_media_name = gr.Textbox(label="Custom Chat Name (optional)", visible=True) + save_chat_history_to_db = gr.Button("Save Chat History to Database") + save_status = gr.Textbox(label="Save Status", interactive=False) + + with gr.Column(scale=2): + chat_history = gr.Chatbot(label="Conversation", height=800) + user_input = gr.Textbox(label="Your message") + send_message_button = gr.Button("Send Message") + answer_for_me_button = gr.Button("Answer for Me") + continue_talking_button = gr.Button("Continue Talking") + regenerate_button = gr.Button("Regenerate Last Message") + clear_chat_button = gr.Button("Clear Chat") + save_snapshot_button = gr.Button("Save Chat Snapshot") + update_chat_dropdown = gr.Dropdown(label="Select Chat to Update", choices=[], visible=False) + load_selected_chat_button = gr.Button("Load Selected Chat", visible=False) + update_chat_button = gr.Button("Update Selected Chat", visible=False) + + # States + character_data = gr.State(None) + user_name = gr.State("") + selected_chat_id = gr.State(None) # To track the selected chat for updates + + # Callback Functions + + def search_existing_chats(query): + results, message = search_character_chats(query) + if results: + # Format search results for dropdown + formatted_results = [ + f"{chat['conversation_name']} (ID: {chat['id']})" for chat in results + ] + else: + formatted_results = [] + return formatted_results, message + + def load_selected_chat_from_search(selected_chat, user_name): + if not selected_chat: + return None, [], None, "No chat selected." + + try: + chat_id_match = re.search(r'\(ID:\s*(\d+)\)', selected_chat) + if not chat_id_match: + return None, [], None, "Invalid chat selection format." + + chat_id = int(chat_id_match.group(1)) + + # Use the new function to load chat and character data + char_data, chat_history, img = load_chat_and_character(chat_id, user_name) + + if not char_data: + return None, [], None, "Failed to load character data for the selected chat." + + return char_data, chat_history, img, f"Chat '{selected_chat}' loaded successfully." + except Exception as e: + logging.error(f"Error loading selected chat: {e}") + return None, [], None, f"Error loading chat: {e}" + + + def import_chat_history(file, current_history, char_data, user_name_val): + """ + Imports chat history from a file, replacing '{{user}}' with the actual user name. + + Args: + file (file): The uploaded chat history file. + current_history (list): The current chat history. + char_data (dict): The current character data. + user_name_val (str): The user's name. + + Returns: + tuple: Updated chat history, updated character data, and a status message. + """ + loaded_history, char_name = load_chat_history(file) + if loaded_history is None: + return current_history, char_data, "Failed to load chat history." + + # Replace '{{user}}' in the loaded chat history + loaded_history = replace_user_placeholder(loaded_history, user_name_val) + + # Check if the loaded chat is for the current character + if char_data and char_data.get('name') != char_name: + return current_history, char_data, ( + f"Warning: Loaded chat is for character '{char_name}', " + f"but current character is '{char_data.get('name')}'. Chat not imported." + ) + + # If no character is selected, try to load the character from the chat + if not char_data: + characters = get_character_cards() + character = next((char for char in characters if char['name'] == char_name), None) + if character: + char_data = character + # Replace '{{user}}' in the first_message if necessary + if character.get('first_message'): + character['first_message'] = character['first_message'].replace("{{user}}", + user_name_val if user_name_val else "User") + else: + return current_history, char_data, ( + f"Warning: Character '{char_name}' not found. Please select the character manually." + ) + + return loaded_history, char_data, f"Chat history for '{char_name}' imported successfully." + + def load_character(name): + characters = get_character_cards() + character = next((char for char in characters if char['name'] == name), None) + if character: + first_message = character.get('first_message', "Hello! I'm ready to chat.") + return character, [(None, first_message)] if first_message else [], None + return None, [], None + + def load_character_image(name): + character = next((char for char in get_character_cards() if char['name'] == name), None) + if character and 'image' in character and character['image']: + try: + # Decode the base64 image + image_data = base64.b64decode(character['image']) + # Load as PIL Image + img = Image.open(io.BytesIO(image_data)).convert("RGBA") + return img + except Exception as e: + logging.error(f"Error loading image for character '{name}': {e}") + return None + return None + + def character_chat_wrapper( + message, history, char_data, api_endpoint, api_key, + temperature, user_name_val, auto_save + ): + if not char_data: + return history, "Please select a character first." + + user_name_val = user_name_val or "User" + char_name = char_data.get('name', 'AI Assistant') + + # Prepare the character's background information + char_background = f""" + Name: {char_name} + Description: {char_data.get('description', 'N/A')} + Personality: {char_data.get('personality', 'N/A')} + Scenario: {char_data.get('scenario', 'N/A')} + """ + + # Prepare the system prompt + system_message = f"""You are roleplaying as {char_name}. {char_data.get('system_prompt', '')}""" + + # Prepare chat context + media_content = { + 'id': char_name, + 'title': char_name, + 'content': char_background, + 'description': char_data.get('description', ''), + 'personality': char_data.get('personality', ''), + 'scenario': char_data.get('scenario', '') + } + selected_parts = ['description', 'personality', 'scenario'] + + prompt = char_data.get('post_history_instructions', '') + + # Sanitize and format user message + user_message = sanitize_user_input(message) + user_message = replace_placeholders(user_message, char_name, user_name_val) + full_message = f"{user_name_val}: {user_message}" + + # Generate bot response + bot_message = chat( + full_message, + history, + media_content, + selected_parts, + api_endpoint, + api_key, + prompt, + temperature, + system_message + ) + + # Replace placeholders in bot message + bot_message = replace_placeholders(bot_message, char_name, user_name_val) + + # Update history + history.append((user_message, bot_message)) + + # Auto-save if enabled + save_status = "" + if auto_save: + character_id = char_data.get('id') + if character_id: + conversation_name = f"Auto-saved chat {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" + add_character_chat(character_id, conversation_name, history) + save_status = "Chat auto-saved." + else: + save_status = "Character ID not found; chat not saved." + + return history, save_status + + def save_chat_history_to_db_wrapper( + chat_history, conversation_id, media_content, + chat_media_name, char_data, auto_save + ): + if not char_data or not chat_history: + return "No character or chat history available.", "" + + character_id = char_data.get('id') + if not character_id: + return "Character ID not found.", "" + + conversation_name = chat_media_name or f"Chat {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" + chat_id = add_character_chat(character_id, conversation_name, chat_history) + if chat_id: + return f"Chat saved successfully with ID {chat_id}.", "" + else: + return "Failed to save chat.", "" + + def update_character_info(name): + return load_character_and_image(name, user_name.value) + + def on_character_select(name, user_name_val): + logging.debug(f"Character selected: {name}") + char_data, chat_history, img = load_character_and_image(name, user_name_val) + return char_data, chat_history, img + + def clear_chat_history(char_data, user_name_val): + """ + Clears the chat history and initializes it with the character's first message, + replacing the '{{user}}' placeholder with the actual user name. + + Args: + char_data (dict): The current character data. + user_name_val (str): The user's name. + + Returns: + tuple: Updated chat history and the unchanged char_data. + """ + if char_data and 'first_message' in char_data and char_data['first_message']: + # Replace '{{user}}' in the first_message + first_message = char_data['first_message'].replace("{{user}}", + user_name_val if user_name_val else "User") + # Initialize chat history with the updated first_message + return [(None, first_message)], char_data + else: + # If no first_message is defined, simply clear the chat + return [], char_data + + def regenerate_last_message( + history, char_data, api_endpoint, api_key, + temperature, user_name_val, auto_save + ): + """ + Regenerates the last bot message by removing it and resending the corresponding user message. + + Args: + history (list): The current chat history as a list of tuples (user_message, bot_message). + char_data (dict): The current character data. + api_endpoint (str): The API endpoint to use for the LLM. + api_key (str): The API key for authentication. + temperature (float): The temperature setting for the LLM. + user_name_val (str): The user's name. + auto_save (bool): Flag indicating whether to auto-save the chat. + + Returns: + tuple: Updated chat history and a save status message. + """ + if not history: + return history, "No messages to regenerate." + + last_entry = history[-1] + last_user_message, last_bot_message = last_entry + + # Check if the last bot message exists + if last_bot_message is None: + return history, "The last message is not from the bot." + + # Remove the last bot message + new_history = history[:-1] + + # Resend the last user message to generate a new bot response + if not last_user_message: + return new_history, "No user message to regenerate the bot response." + + # Prepare the character's background information + char_name = char_data.get('name', 'AI Assistant') + char_background = f""" + Name: {char_name} + Description: {char_data.get('description', 'N/A')} + Personality: {char_data.get('personality', 'N/A')} + Scenario: {char_data.get('scenario', 'N/A')} + """ + + # Prepare the system prompt for character impersonation + system_message = f"""You are roleplaying as {char_name}, the character described below. Respond to the user's messages in character, maintaining the personality and background provided. Do not break character or refer to yourself as an AI. Always refer to yourself as "{char_name}" and refer to the user as "{user_name_val}". + + {char_background} + + Additional instructions: {char_data.get('post_history_instructions', '')} + """ + + # Prepare media_content and selected_parts + media_content = { + 'id': char_name, + 'title': char_name, + 'content': char_background, + 'description': char_data.get('description', ''), + 'personality': char_data.get('personality', ''), + 'scenario': char_data.get('scenario', '') + } + selected_parts = ['description', 'personality', 'scenario'] + + prompt = char_data.get('post_history_instructions', '') + + # Prepare the input for the chat function + full_message = f"{user_name_val}: {last_user_message}" if last_user_message else f"{user_name_val}: " + + # Call the chat function to get a new bot message + bot_message = chat( + full_message, + new_history, + media_content, + selected_parts, + api_endpoint, + api_key, + prompt, + temperature, + system_message + ) + + # Append the new bot message to the history + new_history.append((last_user_message, bot_message)) + + # Auto-save if enabled + if auto_save: + character_id = char_data.get('id') + if character_id: + conversation_name = f"Auto-saved chat {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" + add_character_chat(character_id, conversation_name, new_history) + save_status = "Chat auto-saved." + else: + save_status = "Character ID not found; chat not saved." + else: + save_status = "" + + return new_history, save_status + + def toggle_chat_file_upload(): + return gr.update(visible=True) + + def save_untracked_chat_action(history, char_data): + if not char_data or not history: + return "No chat to save or character not selected." + + character_id = char_data.get('id') + if not character_id: + return "Character ID not found." + + conversation_name = f"Snapshot {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" + chat_id = add_character_chat(character_id, conversation_name, history, is_snapshot=True) + if chat_id: + return f"Chat snapshot saved successfully with ID {chat_id}." + else: + return "Failed to save chat snapshot." + + def select_chat_for_update(): + # Fetch all chats for the selected character + if character_data.value: + character_id = character_data.value.get('id') + if character_id: + chats = get_character_chats(character_id) + chat_choices = [ + f"{chat['conversation_name']} (ID: {chat['id']})" for chat in chats + ] + return gr.update(choices=chat_choices), None + return gr.update(choices=[]), "No character selected." + + def load_selected_chat(chat_selection): + if not chat_selection: + return [], "No chat selected." + + try: + chat_id = int(chat_selection.split('(ID: ')[1].rstrip(')')) + chat = get_character_chat_by_id(chat_id) + if chat: + history = chat['chat_history'] + selected_chat_id.value = chat_id # Update the selected_chat_id state + return history, f"Loaded chat '{chat['conversation_name']}' successfully." + else: + return [], "Chat not found." + except Exception as e: + logging.error(f"Error loading selected chat: {e}") + return [], f"Error loading chat: {e}" + + def update_chat(chat_id, updated_history): + success = update_character_chat(chat_id, updated_history) + if success: + return "Chat updated successfully." + else: + return "Failed to update chat." + + def continue_talking( + history, char_data, api_endpoint, api_key, + temperature, user_name_val, auto_save + ): + """ + Causes the character to continue the conversation or think out loud. + """ + if not char_data: + return history, "Please select a character first." + + user_name_val = user_name_val or "User" + char_name = char_data.get('name', 'AI Assistant') + + # Prepare the character's background information + char_background = f""" + Name: {char_name} + Description: {char_data.get('description', 'N/A')} + Personality: {char_data.get('personality', 'N/A')} + Scenario: {char_data.get('scenario', 'N/A')} + """ + + # Prepare the system prompt + system_message = f"""You are roleplaying as {char_name}. {char_data.get('system_prompt', '')} + If the user does not respond, continue expressing your thoughts or continue the conversation by thinking out loud. If thinking out loud, prefix the message with "Thinking: ".""" + + # Prepare chat context + media_content = { + 'id': char_name, + 'title': char_name, + 'content': char_background, + 'description': char_data.get('description', ''), + 'personality': char_data.get('personality', ''), + 'scenario': char_data.get('scenario', '') + } + selected_parts = ['description', 'personality', 'scenario'] + + prompt = char_data.get('post_history_instructions', '') + + # Simulate empty user input + user_message = "" + + # Generate bot response + bot_message = chat( + user_message, + history, + media_content, + selected_parts, + api_endpoint, + api_key, + prompt, + temperature, + system_message + ) + + # Replace placeholders in bot message + bot_message = replace_placeholders(bot_message, char_name, user_name_val) + + # Update history + history.append((None, bot_message)) + + # Auto-save if enabled + save_status = "" + if auto_save: + character_id = char_data.get('id') + if character_id: + conversation_name = f"Auto-saved chat {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" + add_character_chat(character_id, conversation_name, history) + save_status = "Chat auto-saved." + else: + save_status = "Character ID not found; chat not saved." + + return history, save_status + + def answer_for_me( + history, char_data, api_endpoint, api_key, + temperature, user_name_val, auto_save + ): + """ + Generates a likely user response and continues the conversation. + """ + if not char_data: + return history, "Please select a character first." + + user_name_val = user_name_val or "User" + char_name = char_data.get('name', 'AI Assistant') + + # Prepare the character's background information + char_background = f""" + Name: {char_name} + Description: {char_data.get('description', 'N/A')} + Personality: {char_data.get('personality', 'N/A')} + Scenario: {char_data.get('scenario', 'N/A')} + """ + + # Prepare system message for generating user's response + system_message_user = f"""You are simulating the user {user_name_val}. Based on the conversation so far, generate a natural and appropriate response that {user_name_val} might say next. The response should fit the context and flow of the conversation. ONLY SPEAK FOR {user_name_val}.""" + + # Prepare chat context + media_content = { + 'id': char_name, + 'title': char_name, + 'content': char_background, + 'description': char_data.get('description', ''), + 'personality': char_data.get('personality', ''), + 'scenario': char_data.get('scenario', '') + } + selected_parts = ['description', 'personality', 'scenario'] + + # Generate user response + user_response = chat( + "", # No new message + history, + media_content, + selected_parts, + api_endpoint, + api_key, + prompt="", + temperature=temperature, + system_message=system_message_user + ) + + # Append the generated user response to history + history.append((user_response, None)) + + # Now generate the character's response to this user response + # Prepare the system message for the character + system_message_bot = f"""You are roleplaying as {char_name}. {char_data.get('system_prompt', '')}""" + + bot_message = chat( + f"{user_name_val}: {user_response}", + history[:-1], + media_content, + selected_parts, + api_endpoint, + api_key, + prompt=char_data.get('post_history_instructions', ''), + temperature=temperature, + system_message=system_message_bot + ) + + # Replace placeholders in bot message + bot_message = replace_placeholders(bot_message, char_name, user_name_val) + + # Update history with bot's response + history[-1] = (user_response, bot_message) + + # Auto-save if enabled + save_status = "" + if auto_save: + character_id = char_data.get('id') + if character_id: + conversation_name = f"Auto-saved chat {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" + add_character_chat(character_id, conversation_name, history) + save_status = "Chat auto-saved." + else: + save_status = "Character ID not found; chat not saved." + + return history, save_status + + + # Define States for conversation_id and media_content, which are required for saving chat history + conversation_id = gr.State(str(uuid.uuid4())) + media_content = gr.State({}) + + # Button Callbacks + + # Add the new button callbacks here + answer_for_me_button.click( + fn=answer_for_me, + inputs=[ + chat_history, + character_data, + api_name_input, + api_key_input, + temperature_slider, + user_name_input, + auto_save_checkbox + ], + outputs=[chat_history, save_status] + ) + + continue_talking_button.click( + fn=continue_talking, + inputs=[ + chat_history, + character_data, + api_name_input, + api_key_input, + temperature_slider, + user_name_input, + auto_save_checkbox + ], + outputs=[chat_history, save_status] + ) + + import_card_button.click( + fn=import_character_card, + inputs=[character_card_upload], + outputs=[character_data, character_dropdown, save_status] + ) + + load_characters_button.click( + fn=lambda: gr.update(choices=[f"{char['name']} (ID: {char['id']})" for char in get_character_cards()]), + outputs=character_dropdown + ) + + # FIXME user_name_val = validate_user_name(user_name_val) + clear_chat_button.click( + fn=clear_chat_history, + inputs=[character_data, user_name_input], + outputs=[chat_history, character_data] + ) + + character_dropdown.change( + fn=extract_character_id, + inputs=[character_dropdown], + outputs=character_data + ).then( + fn=load_character_wrapper, + inputs=[character_data, user_name_input], + outputs=[character_data, chat_history, character_image] + ) + + send_message_button.click( + fn=character_chat_wrapper, + inputs=[ + user_input, + chat_history, + character_data, + api_name_input, + api_key_input, + temperature_slider, + user_name_input, + auto_save_checkbox + ], + outputs=[chat_history, save_status] + ).then(lambda: "", outputs=user_input) + + regenerate_button.click( + fn=regenerate_last_message, + inputs=[ + chat_history, + character_data, + api_name_input, + api_key_input, + temperature_slider, + user_name_input, + auto_save_checkbox + ], + outputs=[chat_history, save_status] + ) + + import_chat_button.click( + fn=lambda: gr.update(visible=True), + outputs=chat_file_upload + ) + + chat_file_upload.change( + fn=import_chat_history, + inputs=[chat_file_upload, chat_history, character_data], + outputs=[chat_history, character_data, save_status] + ) + + save_chat_history_to_db.click( + fn=save_chat_history_to_db_wrapper, + inputs=[ + chat_history, + conversation_id, + media_content, + chat_media_name, + character_data, + auto_save_checkbox # Pass the auto_save state + ], + outputs=[conversation_id, save_status] + ) + + # Populate the update_chat_dropdown based on selected character + character_dropdown.change( + fn=select_chat_for_update, + inputs=[], + outputs=[update_chat_dropdown, save_status] + ) + + load_selected_chat_button.click( + fn=load_selected_chat, + inputs=[update_chat_dropdown], + outputs=[chat_history, save_status] + ) + + save_snapshot_button.click( + fn=save_untracked_chat_action, + inputs=[chat_history, character_data], + outputs=save_status + ) + + update_chat_button.click( + fn=update_chat, + inputs=[selected_chat_id, chat_history], + outputs=save_status + ) + + # Search Chats + chat_search_button.click( + fn=search_existing_chats, + inputs=[chat_search_query], + outputs=[chat_search_dropdown, save_status] + ).then( + fn=lambda choices, msg: gr.update(choices=choices, visible=True) if choices else gr.update(visible=False), + inputs=[chat_search_dropdown, save_status], + outputs=[chat_search_dropdown] + ) + + # Load Selected Chat from Search + load_chat_button.click( + fn=load_selected_chat_from_search, + inputs=[chat_search_dropdown, user_name_input], + outputs=[character_data, chat_history, character_image, save_status] + ) + + # Show Load Chat Button when a chat is selected + chat_search_dropdown.change( + fn=lambda selected: gr.update(visible=True) if selected else gr.update(visible=False), + inputs=[chat_search_dropdown], + outputs=[load_chat_button] + ) + + + return character_data, chat_history, user_input, user_name, character_image + + +def create_character_chat_mgmt_tab(): + with gr.TabItem("Character and Chat Management", visible=True): + gr.Markdown("# Character and Chat Management") + + with gr.Row(): + # Left Column: Character Import and Chat Management + with gr.Column(scale=1): + gr.Markdown("## Import Characters") + character_files = gr.File( + label="Upload Character Files (PNG, WEBP, JSON)", + file_types=[".png", ".webp", ".json"], + file_count="multiple" + ) + import_characters_button = gr.Button("Import Characters") + import_status = gr.Markdown("") + + # Right Column: Character Selection and Image Display + with gr.Column(scale=2): + gr.Markdown("## Select Character") + characters = get_character_cards() + character_choices = [f"{char['name']} (ID: {char['id']})" for char in characters] + load_characters_button = gr.Button("Load Existing Characters") + select_character = gr.Dropdown(label="Select Character", choices=character_choices, interactive=True) + character_image = gr.Image(label="Character Image", type="pil", interactive=False) + + gr.Markdown("## Search Conversations") + search_query = gr.Textbox(label="Search Conversations", placeholder="Enter search keywords") + search_button = gr.Button("Search") + search_results = gr.Dropdown(label="Search Results", choices=[], visible=False) + search_status = gr.Markdown("", visible=True) + + with gr.Row(): + gr.Markdown("## Chat Management") + select_chat = gr.Dropdown(label="Select Chat", choices=[], visible=False, interactive=True) + load_chat_button = gr.Button("Load Selected Chat", visible=False) + conversation_list = gr.Dropdown(label="Select Conversation or Character", choices=[]) + conversation_mapping = gr.State({}) + + with gr.Tabs(): + with gr.TabItem("Edit", visible=True): + chat_content = gr.TextArea(label="Chat/Character Content (JSON)", lines=20, max_lines=50) + save_button = gr.Button("Save Changes") + delete_button = gr.Button("Delete Conversation/Character", variant="stop") + + with gr.TabItem("Preview", visible=True): + chat_preview = gr.HTML(label="Chat/Character Preview") + result_message = gr.Markdown("") + + # Callback Functions + + def load_character_image(character_selection): + if not character_selection: + return None + + try: + character_id = int(character_selection.split('(ID: ')[1].rstrip(')')) + character = get_character_card_by_id(character_id) + if character and 'image' in character: + image_data = base64.b64decode(character['image']) + img = Image.open(io.BytesIO(image_data)) + return img + except Exception as e: + logging.error(f"Error loading character image: {e}") + + return None + + def search_conversations_or_characters(query, selected_character): + if not query.strip(): + return gr.update(choices=[], visible=False), "Please enter a search query." + + try: + # Extract character ID from the selected character + character_id = None + if selected_character: + character_id = int(selected_character.split('(ID: ')[1].rstrip(')')) + + # Search Chats using FTS5, filtered by character_id if provided + chat_results, chat_message = search_character_chats(query, character_id) + + # Format chat results + formatted_chat_results = [ + f"Chat: {chat['conversation_name']} (ID: {chat['id']})" for chat in chat_results + ] + + # If no character is selected, also search for characters + if not character_id: + characters = get_character_cards() + filtered_characters = [ + char for char in characters + if query.lower() in char['name'].lower() + ] + formatted_character_results = [ + f"Character: {char['name']} (ID: {char['id']})" for char in filtered_characters + ] + else: + formatted_character_results = [] + + # Combine results + all_choices = formatted_chat_results + formatted_character_results + + if all_choices: + return gr.update(choices=all_choices, visible=True), chat_message + else: + return gr.update(choices=[], visible=False), f"No results found for '{query}'." + + except Exception as e: + logging.error(f"Error during search: {e}") + return gr.update(choices=[], visible=False), f"Error occurred during search: {e}" + + def load_conversation_or_character(selected, conversation_mapping): + if not selected or selected not in conversation_mapping: + return "", "

No selection made.

" + + selected_id = conversation_mapping[selected] + if selected.startswith("Chat:"): + chat = get_character_chat_by_id(selected_id) + if chat: + json_content = json.dumps({ + "conversation_id": chat['id'], + "conversation_name": chat['conversation_name'], + "messages": chat['chat_history'] + }, indent=2) + + html_preview = create_chat_preview_html(chat['chat_history']) + return json_content, html_preview + elif selected.startswith("Character:"): + character = get_character_card_by_id(selected_id) + if character: + json_content = json.dumps({ + "id": character['id'], + "name": character['name'], + "description": character['description'], + "personality": character['personality'], + "scenario": character['scenario'], + "post_history_instructions": character['post_history_instructions'], + "first_mes": character['first_mes'], + "mes_example": character['mes_example'], + "creator_notes": character.get('creator_notes', ''), + "system_prompt": character.get('system_prompt', ''), + "tags": character.get('tags', []), + "creator": character.get('creator', ''), + "character_version": character.get('character_version', ''), + "extensions": character.get('extensions', {}) + }, indent=2) + + html_preview = create_character_preview_html(character) + return json_content, html_preview + + return "", "

Unable to load the selected item.

" + + def validate_content(selected, content): + try: + data = json.loads(content) + if selected.startswith("Chat:"): + assert "conversation_id" in data and "messages" in data + elif selected.startswith("Character:"): + assert "id" in data and "name" in data + return True, data + except Exception as e: + return False, f"Invalid JSON: {e}" + + def save_conversation_or_character(selected, conversation_mapping, content): + if not selected or selected not in conversation_mapping: + return "Please select an item to save.", "

No changes made.

" + + is_valid, result = validate_content(selected, content) + if not is_valid: + return f"Error: {result}", "

No changes made due to validation error.

" + + selected_id = conversation_mapping[selected] + + if selected.startswith("Chat:"): + success = update_character_chat(selected_id, result['messages']) + return ("Chat updated successfully." if success else "Failed to update chat."), ("

Chat updated.

" if success else "

Failed to update chat.

") + elif selected.startswith("Character:"): + success = update_character_card(selected_id, result) + return ("Character updated successfully." if success else "Failed to update character."), ("

Character updated.

" if success else "

Failed to update character.

") + + return "Unknown item type.", "

No changes made.

" + + def delete_conversation_or_character(selected, conversation_mapping): + if not selected or selected not in conversation_mapping: + return "Please select an item to delete.", "

No changes made.

", gr.update(choices=[]) + + selected_id = conversation_mapping[selected] + + if selected.startswith("Chat:"): + success = delete_character_chat(selected_id) + elif selected.startswith("Character:"): + success = delete_character_card(selected_id) + else: + return "Unknown item type.", "

No changes made.

", gr.update() + + if success: + updated_choices = [choice for choice in conversation_mapping.keys() if choice != selected] + conversation_mapping.value.pop(selected, None) + return f"{selected.split(':')[0]} deleted successfully.", f"

{selected.split(':')[0]} deleted.

", gr.update(choices=updated_choices) + else: + return f"Failed to delete {selected.split(':')[0].lower()}.", f"

Failed to delete {selected.split(':')[0].lower()}.

", gr.update() + + def populate_chats(character_selection): + if not character_selection: + return gr.update(choices=[], visible=False), "Please select a character first." + + try: + character_id = int(character_selection.split('(ID: ')[1].rstrip(')')) + chats = get_character_chats(character_id=character_id) + + if not chats: + return gr.update(choices=[], visible=False), f"No chats found for the selected character." + + formatted_chats = [f"{chat['conversation_name']} (ID: {chat['id']})" for chat in chats] + return gr.update(choices=formatted_chats, visible=True), f"Found {len(formatted_chats)} chat(s)." + except Exception as e: + logging.error(f"Error populating chats: {e}") + return gr.update(choices=[], visible=False), f"Error occurred: {e}" + + def load_chat_from_character(selected_chat): + if not selected_chat: + return "", "

No chat selected.

" + + try: + chat_id = int(selected_chat.split('(ID: ')[1].rstrip(')')) + chat = get_character_chat_by_id(chat_id) + if not chat: + return "", "

Selected chat not found.

" + + json_content = json.dumps({ + "conversation_id": chat['id'], + "conversation_name": chat['conversation_name'], + "messages": chat['chat_history'] + }, indent=2) + + html_preview = create_chat_preview_html(chat['chat_history']) + return json_content, html_preview + except Exception as e: + logging.error(f"Error loading chat: {e}") + return "", f"

Error loading chat: {e}

" + + def create_chat_preview_html(chat_history): + html_preview = "
" + for user_msg, bot_msg in chat_history: + user_style = "background-color: #e6f3ff; padding: 10px; border-radius: 5px; margin-bottom: 5px;" + bot_style = "background-color: #f0f0f0; padding: 10px; border-radius: 5px; margin-bottom: 10px;" + html_preview += f"
User: {user_msg}
" + html_preview += f"
Bot: {bot_msg}
" + html_preview += "
" + return html_preview + + def create_character_preview_html(character): + return f""" +
+

{character['name']}

+

Description: {character['description']}

+

Personality: {character['personality']}

+

Scenario: {character['scenario']}

+

First Message: {character['first_mes']}

+

Example Message: {character['mes_example']}

+

Post History Instructions: {character['post_history_instructions']}

+

System Prompt: {character.get('system_prompt', 'N/A')}

+

Tags: {', '.join(character.get('tags', []))}

+

Creator: {character.get('creator', 'N/A')}

+

Version: {character.get('character_version', 'N/A')}

+
+ """ + def import_multiple_characters(files): + if not files: + return "No files provided for character import." + + results = [] + for file in files: + result, _, message = import_character_card(file) + if result: + results.append(f"Imported: {result['name']}") + else: + results.append(f"Failed: {file.name} - {message}") + + # Refresh character choices + characters = get_character_cards() + character_choices = [f"{char['name']} (ID: {char['id']})" for char in characters] + select_character.choices = character_choices + + return "Import results:\n" + "\n".join(results) + + # Register new callback for character import + import_characters_button.click( + fn=import_multiple_characters, + inputs=[character_files], + outputs=[import_status] + ).then( + fn=lambda: gr.update(choices=[f"{char['name']} (ID: {char['id']})" for char in get_character_cards()]), + outputs=select_character + ) + + # Register Callback Functions with Gradio Components + search_button.click( + fn=search_conversations_or_characters, + inputs=[search_query, select_character], + outputs=[search_results, search_status] + ) + + search_results.change( + fn=load_conversation_or_character, + inputs=[search_results, conversation_mapping], + outputs=[chat_content, chat_preview] + ) + + save_button.click( + fn=save_conversation_or_character, + inputs=[conversation_list, conversation_mapping, chat_content], + outputs=[result_message, chat_preview] + ) + + delete_button.click( + fn=delete_conversation_or_character, + inputs=[conversation_list, conversation_mapping], + outputs=[result_message, chat_preview, conversation_list] + ) + + select_character.change( + fn=load_character_image, + inputs=[select_character], + outputs=[character_image] + ).then( + fn=populate_chats, + inputs=[select_character], + outputs=[select_chat, search_status] + ) + + select_chat.change( + fn=load_chat_from_character, + inputs=[select_chat], + outputs=[chat_content, chat_preview] + ) + + load_chat_button.click( + fn=load_chat_from_character, + inputs=[select_chat], + outputs=[chat_content, chat_preview] + ) + + load_characters_button.click( + fn=lambda: gr.update(choices=[f"{char['name']} (ID: {char['id']})" for char in get_character_cards()]), + outputs=select_character + ) + + return ( + character_files, import_characters_button, import_status, + search_query, search_button, search_results, search_status, + select_character, select_chat, load_chat_button, + conversation_list, conversation_mapping, + chat_content, save_button, delete_button, + chat_preview, result_message, character_image + ) + +def create_custom_character_card_tab(): + with gr.TabItem("Create a New Character Card", visible=True): + gr.Markdown("# Create a New Character Card (v2)") + + with gr.Row(): + with gr.Column(): + # Input fields for character card data + name_input = gr.Textbox(label="Name", placeholder="Enter character name") + description_input = gr.TextArea(label="Description", placeholder="Enter character description") + personality_input = gr.TextArea(label="Personality", placeholder="Enter character personality") + scenario_input = gr.TextArea(label="Scenario", placeholder="Enter character scenario") + first_mes_input = gr.TextArea(label="First Message", placeholder="Enter the first message") + mes_example_input = gr.TextArea(label="Example Messages", placeholder="Enter example messages") + creator_notes_input = gr.TextArea(label="Creator Notes", placeholder="Enter notes for the creator") + system_prompt_input = gr.TextArea(label="System Prompt", placeholder="Enter system prompt") + post_history_instructions_input = gr.TextArea(label="Post History Instructions", placeholder="Enter post history instructions") + alternate_greetings_input = gr.TextArea( + label="Alternate Greetings (one per line)", + placeholder="Enter alternate greetings, one per line" + ) + tags_input = gr.Textbox(label="Tags", placeholder="Enter tags, separated by commas") + creator_input = gr.Textbox(label="Creator", placeholder="Enter creator name") + character_version_input = gr.Textbox(label="Character Version", placeholder="Enter character version") + extensions_input = gr.TextArea( + label="Extensions (JSON)", + placeholder="Enter extensions as JSON (optional)" + ) + image_input = gr.Image(label="Character Image", type="pil") + + # Buttons + save_button = gr.Button("Save Character Card") + download_button = gr.Button("Download Character Card") + download_image_button = gr.Button("Download Character Card as Image") + + # Output status and outputs + save_status = gr.Markdown("") + download_output = gr.File(label="Download Character Card", interactive=False) + download_image_output = gr.File(label="Download Character Card as Image", interactive=False) + + # Import PngInfo + from PIL.PngImagePlugin import PngInfo + + # Callback Functions + def build_character_card( + name, description, personality, scenario, first_mes, mes_example, + creator_notes, system_prompt, post_history_instructions, + alternate_greetings_str, tags_str, creator, character_version, + extensions_str + ): + # Parse alternate_greetings from multiline string + alternate_greetings = [line.strip() for line in alternate_greetings_str.strip().split('\n') if line.strip()] + + # Parse tags from comma-separated string + tags = [tag.strip() for tag in tags_str.strip().split(',') if tag.strip()] + + # Parse extensions from JSON string + try: + extensions = json.loads(extensions_str) if extensions_str.strip() else {} + except json.JSONDecodeError as e: + extensions = {} + logging.error(f"Error parsing extensions JSON: {e}") + + # Build the character card dictionary according to V2 spec + character_card = { + 'spec': 'chara_card_v2', + 'spec_version': '2.0', + 'data': { + 'name': name, + 'description': description, + 'personality': personality, + 'scenario': scenario, + 'first_mes': first_mes, + 'mes_example': mes_example, + 'creator_notes': creator_notes, + 'system_prompt': system_prompt, + 'post_history_instructions': post_history_instructions, + 'alternate_greetings': alternate_greetings, + 'tags': tags, + 'creator': creator, + 'character_version': character_version, + 'extensions': extensions, + } + } + return character_card + + def validate_character_card_data(character_card): + """ + Validates the character card data using the extended validation logic. + """ + is_valid, validation_messages = validate_v2_card(character_card) + return is_valid, validation_messages + + def save_character_card( + name, description, personality, scenario, first_mes, mes_example, + creator_notes, system_prompt, post_history_instructions, + alternate_greetings_str, tags_str, creator, character_version, + extensions_str, image + ): + # Build the character card + character_card = build_character_card( + name, description, personality, scenario, first_mes, mes_example, + creator_notes, system_prompt, post_history_instructions, + alternate_greetings_str, tags_str, creator, character_version, + extensions_str + ) + + # Validate the character card + is_valid, validation_messages = validate_character_card_data(character_card) + if not is_valid: + # Return validation errors + validation_output = "Character card validation failed:\n" + validation_output += "\n".join(validation_messages) + return validation_output + + # If image is provided, encode it to base64 + if image: + img_byte_arr = io.BytesIO() + image.save(img_byte_arr, format='PNG') + character_card['data']['image'] = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8') + + # Save character card to database + character_id = add_character_card(character_card['data']) + if character_id: + return f"Character card '{name}' saved successfully." + else: + return f"Failed to save character card '{name}'. It may already exist." + + def download_character_card( + name, description, personality, scenario, first_mes, mes_example, + creator_notes, system_prompt, post_history_instructions, + alternate_greetings_str, tags_str, creator, character_version, + extensions_str, image + ): + # Build the character card + character_card = build_character_card( + name, description, personality, scenario, first_mes, mes_example, + creator_notes, system_prompt, post_history_instructions, + alternate_greetings_str, tags_str, creator, character_version, + extensions_str + ) + + # Validate the character card + is_valid, validation_messages = validate_character_card_data(character_card) + if not is_valid: + # Return validation errors + validation_output = "Character card validation failed:\n" + validation_output += "\n".join(validation_messages) + return gr.update(value=None), validation_output # Return None for the file output + + # If image is provided, include it as base64 + if image: + img_byte_arr = io.BytesIO() + image.save(img_byte_arr, format='PNG') + character_card['data']['image'] = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8') + + # Convert to JSON string + json_str = json.dumps(character_card, indent=2) + + # Write the JSON to a temporary file + with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', encoding='utf-8') as temp_file: + temp_file.write(json_str) + temp_file_path = temp_file.name + + # Return the file path and clear validation output + return temp_file_path, "" + + def download_character_card_as_image( + name, description, personality, scenario, first_mes, mes_example, + creator_notes, system_prompt, post_history_instructions, + alternate_greetings_str, tags_str, creator, character_version, + extensions_str, image + ): + # Build the character card + character_card = build_character_card( + name, description, personality, scenario, first_mes, mes_example, + creator_notes, system_prompt, post_history_instructions, + alternate_greetings_str, tags_str, creator, character_version, + extensions_str + ) + + # Validate the character card + is_valid, validation_messages = validate_character_card_data(character_card) + if not is_valid: + # Return validation errors + validation_output = "Character card validation failed:\n" + validation_output += "\n".join(validation_messages) + return gr.update(value=None), validation_output # Return None for the file output + + # Convert the character card JSON to a string + json_str = json.dumps(character_card, indent=2) + + # Encode the JSON string to base64 + chara_content = base64.b64encode(json_str.encode('utf-8')).decode('utf-8') + + # Create PNGInfo object to hold metadata + png_info = PngInfo() + png_info.add_text('chara', chara_content) + + # If image is provided, use it; otherwise, create a blank image + if image: + img = image.copy() + else: + # Create a default blank image + img = Image.new('RGB', (512, 512), color='white') + + # Save the image to a temporary file with metadata + with tempfile.NamedTemporaryFile(mode='wb', delete=False, suffix='.png') as temp_file: + img.save(temp_file, format='PNG', pnginfo=png_info) + temp_file_path = temp_file.name + + # Return the file path and clear validation output + return temp_file_path, "" + + # Include the validate_v2_card function here (from previous code) + + # Button Callbacks + save_button.click( + fn=save_character_card, + inputs=[ + name_input, description_input, personality_input, scenario_input, + first_mes_input, mes_example_input, creator_notes_input, system_prompt_input, + post_history_instructions_input, alternate_greetings_input, tags_input, + creator_input, character_version_input, extensions_input, image_input + ], + outputs=[save_status] + ) + + download_button.click( + fn=download_character_card, + inputs=[ + name_input, description_input, personality_input, scenario_input, + first_mes_input, mes_example_input, creator_notes_input, system_prompt_input, + post_history_instructions_input, alternate_greetings_input, tags_input, + creator_input, character_version_input, extensions_input, image_input + ], + outputs=[download_output, save_status] + ) + + download_image_button.click( + fn=download_character_card_as_image, + inputs=[ + name_input, description_input, personality_input, scenario_input, + first_mes_input, mes_example_input, creator_notes_input, system_prompt_input, + post_history_instructions_input, alternate_greetings_input, tags_input, + creator_input, character_version_input, extensions_input, image_input + ], + outputs=[download_image_output, save_status] + ) + + +def create_character_card_validation_tab(): + with gr.TabItem("Validate Character Card", visible=True): + gr.Markdown("# Validate Character Card (v2)") + gr.Markdown("Upload a character card (PNG, WEBP, or JSON) to validate whether it conforms to the Character Card V2 specification.") + + with gr.Row(): + with gr.Column(): + # File uploader + file_upload = gr.File( + label="Upload Character Card (PNG, WEBP, JSON)", + file_types=[".png", ".webp", ".json"] + ) + # Validation button + validate_button = gr.Button("Validate Character Card") + # Output area for validation results + validation_output = gr.Markdown("") + + # Callback Functions + def validate_character_card(file): + if file is None: + return "No file provided for validation." + + try: + if file.name.lower().endswith(('.png', '.webp')): + json_data = extract_json_from_image(file) + if not json_data: + return "Failed to extract JSON data from the image. The image might not contain embedded character card data." + elif file.name.lower().endswith('.json'): + with open(file.name, 'r', encoding='utf-8') as f: + json_data = f.read() + else: + return "Unsupported file type. Please upload a PNG, WEBP, or JSON file." + + # Parse the JSON content + try: + card_data = json.loads(json_data) + except json.JSONDecodeError as e: + return f"JSON decoding error: {e}" + + # Validate the character card + is_valid, validation_messages = validate_v2_card(card_data) + + # Prepare the validation output + if is_valid: + return "Character card is valid according to the V2 specification." + else: + # Concatenate all validation error messages + validation_output = "Character card validation failed:\n" + validation_output += "\n".join(validation_messages) + return validation_output + + except Exception as e: + logging.error(f"Error validating character card: {e}") + return f"An unexpected error occurred during validation: {e}" + + def validate_v2_card(card_data): + """ + Validate a character card according to the V2 specification. + + Args: + card_data (dict): The parsed character card data. + + Returns: + Tuple[bool, List[str]]: A tuple containing a boolean indicating validity and a list of validation messages. + """ + validation_messages = [] + + # Check top-level fields + if 'spec' not in card_data: + validation_messages.append("Missing 'spec' field.") + elif card_data['spec'] != 'chara_card_v2': + validation_messages.append(f"Invalid 'spec' value: {card_data['spec']}. Expected 'chara_card_v2'.") + + if 'spec_version' not in card_data: + validation_messages.append("Missing 'spec_version' field.") + else: + # Ensure 'spec_version' is '2.0' or higher + try: + spec_version = float(card_data['spec_version']) + if spec_version < 2.0: + validation_messages.append(f"'spec_version' must be '2.0' or higher. Found '{card_data['spec_version']}'.") + except ValueError: + validation_messages.append(f"Invalid 'spec_version' format: {card_data['spec_version']}. Must be a number as a string.") + + if 'data' not in card_data: + validation_messages.append("Missing 'data' field.") + return False, validation_messages # Cannot proceed without 'data' field + + data = card_data['data'] + + # Required fields in 'data' + required_fields = ['name', 'description', 'personality', 'scenario', 'first_mes', 'mes_example'] + for field in required_fields: + if field not in data: + validation_messages.append(f"Missing required field in 'data': '{field}'.") + elif not isinstance(data[field], str): + validation_messages.append(f"Field '{field}' must be a string.") + elif not data[field].strip(): + validation_messages.append(f"Field '{field}' cannot be empty.") + + # Optional fields with expected types + optional_fields = { + 'creator_notes': str, + 'system_prompt': str, + 'post_history_instructions': str, + 'alternate_greetings': list, + 'tags': list, + 'creator': str, + 'character_version': str, + 'extensions': dict, + 'character_book': dict # If present, should be a dict + } + + for field, expected_type in optional_fields.items(): + if field in data: + if not isinstance(data[field], expected_type): + validation_messages.append(f"Field '{field}' must be of type '{expected_type.__name__}'.") + elif field == 'extensions': + # Validate that extensions keys are properly namespaced + for key in data[field].keys(): + if '/' not in key and '_' not in key: + validation_messages.append(f"Extension key '{key}' in 'extensions' should be namespaced to prevent conflicts.") + + # If 'alternate_greetings' is present, check that it's a list of non-empty strings + if 'alternate_greetings' in data and isinstance(data['alternate_greetings'], list): + for idx, greeting in enumerate(data['alternate_greetings']): + if not isinstance(greeting, str) or not greeting.strip(): + validation_messages.append(f"Element {idx} in 'alternate_greetings' must be a non-empty string.") + + # If 'tags' is present, check that it's a list of non-empty strings + if 'tags' in data and isinstance(data['tags'], list): + for idx, tag in enumerate(data['tags']): + if not isinstance(tag, str) or not tag.strip(): + validation_messages.append(f"Element {idx} in 'tags' must be a non-empty string.") + + # Validate 'extensions' field + if 'extensions' in data and not isinstance(data['extensions'], dict): + validation_messages.append("Field 'extensions' must be a dictionary.") + + # Validate 'character_book' if present + if 'character_book' in data: + is_valid_book, book_messages = validate_character_book(data['character_book']) + if not is_valid_book: + validation_messages.extend(book_messages) + + is_valid = len(validation_messages) == 0 + return is_valid, validation_messages + + # Button Callback + validate_button.click( + fn=validate_character_card, + inputs=[file_upload], + outputs=[validation_output] + ) + + +def create_export_characters_tab(): + with gr.TabItem("Export Characters", visible=True): + gr.Markdown("# Export Characters") + gr.Markdown("Export character cards individually as JSON files or all together as a ZIP file.") + + with gr.Row(): + with gr.Column(scale=1): + # Dropdown to select a character for individual export + characters = get_character_cards() + character_choices = [f"{char['name']} (ID: {char['id']})" for char in characters] + export_character_dropdown = gr.Dropdown( + label="Select Character to Export", + choices=character_choices + ) + load_characters_button = gr.Button("Load Existing Characters") + export_single_button = gr.Button("Export Selected Character") + export_all_button = gr.Button("Export All Characters") + + with gr.Column(scale=1): + # Output components + export_output = gr.File(label="Exported Character(s)", interactive=False) + export_status = gr.Markdown("") + +# FIXME + def export_single_character_wrapper(character_selection): + file_path, status_message = export_single_character(character_selection) + if file_path: + return gr.update(value=file_path), status_message + else: + return gr.update(value=None), status_message + + def export_all_characters_wrapper(): + zip_path = export_all_characters_as_zip() + characters = get_character_cards() + exported_characters = [char['name'] for char in characters] + status_message = f"Exported {len(exported_characters)} characters successfully:\n" + "\n".join(exported_characters) + return gr.update(value=zip_path), status_message + + # Event listeners + load_characters_button.click( + fn=lambda: gr.update(choices=[f"{char['name']} (ID: {char['id']})" for char in get_character_cards()]), + outputs=export_character_dropdown + ) + + export_single_button.click( + fn=export_single_character_wrapper, + inputs=[export_character_dropdown], + outputs=[export_output, export_status] + ) + + export_all_button.click( + fn=export_all_characters_wrapper, + inputs=[], + outputs=[export_output, export_status] + ) + + return export_character_dropdown, load_characters_button, export_single_button, export_all_button, export_output, export_status + +# +# End of Character_Chat_tab.py +####################################################################################################################### \ No newline at end of file diff --git a/App_Function_Libraries/Gradio_UI/Character_interaction_tab.py b/App_Function_Libraries/Gradio_UI/Character_interaction_tab.py new file mode 100644 index 0000000000000000000000000000000000000000..5d1738052b94369997ea157b13ba718f34b01ed8 --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Character_interaction_tab.py @@ -0,0 +1,511 @@ +# Character_Interaction_tab.py +# Description: This file contains the functions that are used for Character Interactions in the Gradio UI. +# +# Imports +import base64 +import io +import uuid +from datetime import datetime as datetime +import logging +import json +import os +from typing import List, Dict, Tuple, Union + +# +# External Imports +import gradio as gr +from PIL import Image +# +# Local Imports +from App_Function_Libraries.Chat import chat, load_characters, save_chat_history_to_db_wrapper +from App_Function_Libraries.Gradio_UI.Chat_ui import chat_wrapper +from App_Function_Libraries.Gradio_UI.Writing_tab import generate_writing_feedback +# +######################################################################################################################## +# +# Single-Character chat Functions: +# FIXME - add these functions to the Personas library + +def chat_with_character(user_message, history, char_data, api_name_input, api_key): + if char_data is None: + return history, "Please import a character card first." + + bot_message = generate_writing_feedback(user_message, char_data['name'], "Overall", api_name_input, + api_key) + history.append((user_message, bot_message)) + return history, "" + + +def import_character_card(file): + if file is None: + logging.warning("No file provided for character card import") + return None + try: + if file.name.lower().endswith(('.png', '.webp')): + logging.info(f"Attempting to import character card from image: {file.name}") + json_data = extract_json_from_image(file) + if json_data: + logging.info("JSON data extracted from image, attempting to parse") + card_data = import_character_card_json(json_data) + if card_data: + # Save the image data + with Image.open(file) as img: + img_byte_arr = io.BytesIO() + img.save(img_byte_arr, format='PNG') + card_data['image'] = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8') + return card_data + else: + logging.warning("No JSON data found in the image") + else: + logging.info(f"Attempting to import character card from JSON file: {file.name}") + content = file.read().decode('utf-8') + return import_character_card_json(content) + except Exception as e: + logging.error(f"Error importing character card: {e}") + return None + + +def import_character_card_json(json_content): + try: + # Remove any leading/trailing whitespace + json_content = json_content.strip() + + # Log the first 100 characters of the content + logging.debug(f"JSON content (first 100 chars): {json_content[:100]}...") + + card_data = json.loads(json_content) + logging.debug(f"Parsed JSON data keys: {list(card_data.keys())}") + if 'spec' in card_data and card_data['spec'] == 'chara_card_v2': + logging.info("Detected V2 character card") + return card_data['data'] + else: + logging.info("Assuming V1 character card") + return card_data + except json.JSONDecodeError as e: + logging.error(f"JSON decode error: {e}") + logging.error(f"Problematic JSON content: {json_content[:500]}...") + except Exception as e: + logging.error(f"Unexpected error parsing JSON: {e}") + return None + + +def extract_json_from_image(image_file): + logging.debug(f"Attempting to extract JSON from image: {image_file.name}") + try: + with Image.open(image_file) as img: + logging.debug("Image opened successfully") + metadata = img.info + if 'chara' in metadata: + logging.debug("Found 'chara' in image metadata") + chara_content = metadata['chara'] + logging.debug(f"Content of 'chara' metadata (first 100 chars): {chara_content[:100]}...") + try: + decoded_content = base64.b64decode(chara_content).decode('utf-8') + logging.debug(f"Decoded content (first 100 chars): {decoded_content[:100]}...") + return decoded_content + except Exception as e: + logging.error(f"Error decoding base64 content: {e}") + + logging.debug("'chara' not found in metadata, checking for base64 encoded data") + raw_data = img.tobytes() + possible_json = raw_data.split(b'{', 1)[-1].rsplit(b'}', 1)[0] + if possible_json: + try: + decoded = base64.b64decode(possible_json).decode('utf-8') + if decoded.startswith('{') and decoded.endswith('}'): + logging.debug("Found and decoded base64 JSON data") + return '{' + decoded + '}' + except Exception as e: + logging.error(f"Error decoding base64 data: {e}") + + logging.warning("No JSON data found in the image") + except Exception as e: + logging.error(f"Error extracting JSON from image: {e}") + return None + + +def load_chat_history(file): + try: + content = file.read().decode('utf-8') + chat_data = json.loads(content) + return chat_data['history'], chat_data['character'] + except Exception as e: + logging.error(f"Error loading chat history: {e}") + return None, None + + +# +# End of X +###################################################################################################################### +# +# Multi-Character Chat Interface + +# FIXME - refactor and move these functions to the Character_Chat library so that it uses the same functions +def character_interaction_setup(): + characters = load_characters() + return characters, [], None, None + + +def extract_character_response(response: Union[str, Tuple]) -> str: + if isinstance(response, tuple): + # If it's a tuple, try to extract the first string element + for item in response: + if isinstance(item, str): + return item.strip() + # If no string found, return a default message + return "I'm not sure how to respond." + elif isinstance(response, str): + # If it's already a string, just return it + return response.strip() + else: + # For any other type, return a default message + return "I'm having trouble forming a response." + +# def process_character_response(response: str) -> str: +# # Remove any leading explanatory text before the first '---' +# parts = response.split('---') +# if len(parts) > 1: +# return '---' + '---'.join(parts[1:]) +# return response.strip() +def process_character_response(response: Union[str, Tuple]) -> str: + if isinstance(response, tuple): + response = ' '.join(str(item) for item in response if isinstance(item, str)) + + if isinstance(response, str): + # Remove any leading explanatory text before the first '---' + parts = response.split('---') + if len(parts) > 1: + return '---' + '---'.join(parts[1:]) + return response.strip() + else: + return "I'm having trouble forming a response." + +def character_turn(characters: Dict, conversation: List[Tuple[str, str]], + current_character: str, other_characters: List[str], + api_endpoint: str, api_key: str, temperature: float, + scenario: str = "") -> Tuple[List[Tuple[str, str]], str]: + if not current_character or current_character not in characters: + return conversation, current_character + + if not conversation and scenario: + conversation.append(("Scenario", scenario)) + + current_char = characters[current_character] + other_chars = [characters[char] for char in other_characters if char in characters and char != current_character] + + prompt = f"{current_char['name']}'s personality: {current_char['personality']}\n" + for char in other_chars: + prompt += f"{char['name']}'s personality: {char['personality']}\n" + prompt += "Conversation so far:\n" + "\n".join([f"{sender}: {message}" for sender, message in conversation]) + prompt += f"\n\nHow would {current_char['name']} respond?" + + try: + response = chat_wrapper(prompt, conversation, {}, [], api_endpoint, api_key, "", None, False, temperature, "") + processed_response = process_character_response(response) + conversation.append((current_char['name'], processed_response)) + except Exception as e: + error_message = f"Error generating response: {str(e)}" + conversation.append((current_char['name'], error_message)) + + return conversation, current_character + + +def character_interaction(character1: str, character2: str, api_endpoint: str, api_key: str, + num_turns: int, scenario: str, temperature: float, + user_interjection: str = "") -> List[str]: + characters = load_characters() + char1 = characters[character1] + char2 = characters[character2] + conversation = [] + current_speaker = char1 + other_speaker = char2 + + # Add scenario to the conversation start + if scenario: + conversation.append(f"Scenario: {scenario}") + + for turn in range(num_turns): + # Construct the prompt for the current speaker + prompt = f"{current_speaker['name']}'s personality: {current_speaker['personality']}\n" + prompt += f"{other_speaker['name']}'s personality: {other_speaker['personality']}\n" + prompt += f"Conversation so far:\n" + "\n".join( + [msg if isinstance(msg, str) else f"{msg[0]}: {msg[1]}" for msg in conversation]) + + # Add user interjection if provided + if user_interjection and turn == num_turns // 2: + prompt += f"\n\nUser interjection: {user_interjection}\n" + conversation.append(f"User: {user_interjection}") + + prompt += f"\n\nHow would {current_speaker['name']} respond?" + + # FIXME - figure out why the double print is happening + # Get response from the LLM + response = chat_wrapper(prompt, conversation, {}, [], api_endpoint, api_key, "", None, False, temperature, "") + + # Add the response to the conversation + conversation.append((current_speaker['name'], response)) + + # Switch speakers + current_speaker, other_speaker = other_speaker, current_speaker + + # Convert the conversation to a list of strings for output + return [f"{msg[0]}: {msg[1]}" if isinstance(msg, tuple) else msg for msg in conversation] + + +def create_multiple_character_chat_tab(): + with gr.TabItem("Multi-Character Chat", visible=True): + characters, conversation, current_character, other_character = character_interaction_setup() + + with gr.Blocks() as character_interaction: + gr.Markdown("# Multi-Character Chat") + + with gr.Row(): + num_characters = gr.Dropdown(label="Number of Characters", choices=["2", "3", "4"], value="2") + character_selectors = [gr.Dropdown(label=f"Character {i + 1}", choices=list(characters.keys())) for i in + range(4)] + + api_endpoint = gr.Dropdown(label="API Endpoint", + choices=["Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", "DeepSeek", + "Mistral", + "OpenRouter", "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM", + "ollama", "HuggingFace", + "Custom-OpenAI-API"], + value="HuggingFace") + api_key = gr.Textbox(label="API Key (if required)", type="password") + temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, step=0.1, value=0.7) + scenario = gr.Textbox(label="Scenario (optional)", lines=3) + + chat_display = gr.Chatbot(label="Character Interaction") + current_index = gr.State(0) + + next_turn_btn = gr.Button("Next Turn") + narrator_input = gr.Textbox(label="Narrator Input", placeholder="Add a narration or description...") + add_narration_btn = gr.Button("Add Narration") + error_box = gr.Textbox(label="Error Messages", visible=False) + reset_btn = gr.Button("Reset Conversation") + chat_media_name = gr.Textbox(label="Custom Chat Name(optional)", visible=True) + save_chat_history_to_db = gr.Button("Save Chat History to DataBase") + + def update_character_selectors(num): + return [gr.update(visible=True) if i < int(num) else gr.update(visible=False) for i in range(4)] + + num_characters.change( + update_character_selectors, + inputs=[num_characters], + outputs=character_selectors + ) + + def reset_conversation(): + return [], 0, gr.update(value=""), gr.update(value="") + + def take_turn(conversation, current_index, char1, char2, char3, char4, api_endpoint, api_key, temperature, + scenario): + char_selectors = [char for char in [char1, char2, char3, char4] if char] # Remove None values + num_chars = len(char_selectors) + + if num_chars == 0: + return conversation, current_index # No characters selected, return without changes + + if not conversation: + conversation = [] + if scenario: + conversation.append(("Scenario", scenario)) + + current_character = char_selectors[current_index % num_chars] + next_index = (current_index + 1) % num_chars + + prompt = f"Character speaking: {current_character}\nOther characters: {', '.join(char for char in char_selectors if char != current_character)}\n" + prompt += "Generate the next part of the conversation, including character dialogues and actions. Characters should speak in first person." + + response, new_conversation, _ = chat_wrapper(prompt, conversation, {}, [], api_endpoint, api_key, "", + None, False, temperature, "") + + # Format the response + formatted_lines = [] + for line in response.split('\n'): + if ':' in line: + speaker, text = line.split(':', 1) + formatted_lines.append(f"**{speaker.strip()}**: {text.strip()}") + else: + formatted_lines.append(line) + + formatted_response = '\n'.join(formatted_lines) + + # Update the last message in the conversation with the formatted response + if new_conversation: + new_conversation[-1] = (new_conversation[-1][0], formatted_response) + else: + new_conversation.append((current_character, formatted_response)) + + return new_conversation, next_index + + def add_narration(narration, conversation): + if narration: + conversation.append(("Narrator", narration)) + return conversation, "" + + def take_turn_with_error_handling(conversation, current_index, char1, char2, char3, char4, api_endpoint, + api_key, temperature, scenario): + try: + new_conversation, next_index = take_turn(conversation, current_index, char1, char2, char3, char4, + api_endpoint, api_key, temperature, scenario) + return new_conversation, next_index, gr.update(visible=False, value="") + except Exception as e: + error_message = f"An error occurred: {str(e)}" + return conversation, current_index, gr.update(visible=True, value=error_message) + + # Define States for conversation_id and media_content, which are required for saving chat history + media_content = gr.State({}) + conversation_id = gr.State(str(uuid.uuid4())) + + next_turn_btn.click( + take_turn_with_error_handling, + inputs=[chat_display, current_index] + character_selectors + [api_endpoint, api_key, temperature, + scenario], + outputs=[chat_display, current_index, error_box] + ) + + add_narration_btn.click( + add_narration, + inputs=[narrator_input, chat_display], + outputs=[chat_display, narrator_input] + ) + + reset_btn.click( + reset_conversation, + outputs=[chat_display, current_index, scenario, narrator_input] + ) + + # FIXME - Implement saving chat history to database; look at Chat_UI.py for reference + save_chat_history_to_db.click( + save_chat_history_to_db_wrapper, + inputs=[chat_display, conversation_id, media_content, chat_media_name], + outputs=[conversation_id, gr.Textbox(label="Save Status")] + ) + + return character_interaction + +# +# End of Multi-Character chat tab +######################################################################################################################## +# +# Narrator-Controlled Conversation Tab + +# From `Fuzzlewumper` on Reddit. +def create_narrator_controlled_conversation_tab(): + with gr.TabItem("Narrator-Controlled Conversation", visible=True): + gr.Markdown("# Narrator-Controlled Conversation") + + with gr.Row(): + with gr.Column(scale=1): + api_endpoint = gr.Dropdown( + label="API Endpoint", + choices=["Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", "DeepSeek", "Mistral", + "OpenRouter", "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM", "ollama", "HuggingFace", + "Custom-OpenAI-API"], + value="HuggingFace" + ) + api_key = gr.Textbox(label="API Key (if required)", type="password") + temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, step=0.1, value=0.7) + + with gr.Column(scale=2): + narrator_input = gr.Textbox( + label="Narrator Input", + placeholder="Set the scene or provide context...", + lines=3 + ) + + character_inputs = [] + for i in range(4): # Allow up to 4 characters + with gr.Row(): + name = gr.Textbox(label=f"Character {i + 1} Name") + description = gr.Textbox(label=f"Character {i + 1} Description", lines=3) + character_inputs.append((name, description)) + + conversation_display = gr.Chatbot(label="Conversation", height=400) + user_input = gr.Textbox(label="Your Input (optional)", placeholder="Add your own dialogue or action...") + + with gr.Row(): + generate_btn = gr.Button("Generate Next Interaction") + reset_btn = gr.Button("Reset Conversation") + chat_media_name = gr.Textbox(label="Custom Chat Name(optional)", visible=True) + save_chat_history_to_db = gr.Button("Save Chat History to DataBase") + + error_box = gr.Textbox(label="Error Messages", visible=False) + + # Define States for conversation_id and media_content, which are required for saving chat history + conversation_id = gr.State(str(uuid.uuid4())) + media_content = gr.State({}) + + def generate_interaction(conversation, narrator_text, user_text, api_endpoint, api_key, temperature, + *character_data): + try: + characters = [{"name": name.strip(), "description": desc.strip()} + for name, desc in zip(character_data[::2], character_data[1::2]) + if name.strip() and desc.strip()] + + if not characters: + raise ValueError("At least one character must be defined.") + + prompt = f"Narrator: {narrator_text}\n\n" + for char in characters: + prompt += f"Character '{char['name']}': {char['description']}\n" + prompt += "\nGenerate the next part of the conversation, including character dialogues and actions. " + prompt += "Characters should speak in first person. " + if user_text: + prompt += f"\nIncorporate this user input: {user_text}" + prompt += "\nResponse:" + + response, conversation, _ = chat_wrapper(prompt, conversation, {}, [], api_endpoint, api_key, "", None, + False, temperature, "") + + # Format the response + formatted_lines = [] + for line in response.split('\n'): + if ':' in line: + speaker, text = line.split(':', 1) + formatted_lines.append(f"**{speaker.strip()}**: {text.strip()}") + else: + formatted_lines.append(line) + + formatted_response = '\n'.join(formatted_lines) + + # Update the last message in the conversation with the formatted response + if conversation: + conversation[-1] = (conversation[-1][0], formatted_response) + else: + conversation.append((None, formatted_response)) + + return conversation, gr.update(value=""), gr.update(value=""), gr.update(visible=False, value="") + except Exception as e: + error_message = f"An error occurred: {str(e)}" + return conversation, gr.update(), gr.update(), gr.update(visible=True, value=error_message) + + def reset_conversation(): + return [], gr.update(value=""), gr.update(value=""), gr.update(visible=False, value="") + + generate_btn.click( + generate_interaction, + inputs=[conversation_display, narrator_input, user_input, api_endpoint, api_key, temperature] + + [input for char_input in character_inputs for input in char_input], + outputs=[conversation_display, narrator_input, user_input, error_box] + ) + + reset_btn.click( + reset_conversation, + outputs=[conversation_display, narrator_input, user_input, error_box] + ) + + # FIXME - Implement saving chat history to database; look at Chat_UI.py for reference + save_chat_history_to_db.click( + save_chat_history_to_db_wrapper, + inputs=[conversation_display, conversation_id, media_content, chat_media_name], + outputs=[conversation_id, gr.Textbox(label="Save Status")] + ) + + + return api_endpoint, api_key, temperature, narrator_input, conversation_display, user_input, generate_btn, reset_btn, error_box + +# +# End of Narrator-Controlled Conversation tab +######################################################################################################################## \ No newline at end of file diff --git a/App_Function_Libraries/Gradio_UI/Chat_Workflows.py b/App_Function_Libraries/Gradio_UI/Chat_Workflows.py new file mode 100644 index 0000000000000000000000000000000000000000..e47ad4037a40c5e18f57bdd48e3d8c67dab5e1ff --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Chat_Workflows.py @@ -0,0 +1,178 @@ +# Chat_Workflows.py +# Description: UI for Chat Workflows +# +# Imports +import json +import logging +from pathlib import Path +# +# External Imports +import gradio as gr +# +from App_Function_Libraries.Gradio_UI.Chat_ui import chat_wrapper, search_conversations, \ + load_conversation +from App_Function_Libraries.Chat import save_chat_history_to_db_wrapper +# +############################################################################################################ +# +# Functions: + +# Load workflows from a JSON file +json_path = Path('./Helper_Scripts/Workflows/Workflows.json') +with json_path.open('r') as f: + workflows = json.load(f) + + +def chat_workflows_tab(): + with gr.TabItem("Chat Workflows", visible=True): + gr.Markdown("# Workflows using LLMs") + chat_history = gr.State([]) + media_content = gr.State({}) + selected_parts = gr.State([]) + conversation_id = gr.State(None) + workflow_state = gr.State({"current_step": 0, "max_steps": 0, "conversation_id": None}) + + with gr.Row(): + with gr.Column(): + workflow_selector = gr.Dropdown(label="Select Workflow", choices=[wf['name'] for wf in workflows]) + api_selector = gr.Dropdown( + label="Select API Endpoint", + choices=["Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", "DeepSeek", "Mistral", + "OpenRouter", "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM", "ollama", "HuggingFace", + "Custom-OpenAI-API"], + value="HuggingFace" + ) + api_key_input = gr.Textbox(label="API Key (optional)", type="password") + temperature = gr.Slider(label="Temperature", minimum=0.00, maximum=1.0, step=0.05, value=0.7) + save_conversation = gr.Checkbox(label="Save Conversation", value=False) + with gr.Column(): + gr.Markdown("Placeholder") + with gr.Row(): + with gr.Column(): + conversation_search = gr.Textbox(label="Search Conversations") + search_conversations_btn = gr.Button("Search Conversations") + with gr.Column(): + previous_conversations = gr.Dropdown(label="Select Conversation", choices=[], interactive=True) + load_conversations_btn = gr.Button("Load Selected Conversation") + with gr.Row(): + with gr.Column(): + context_input = gr.Textbox(label="Initial Context", lines=5) + chatbot = gr.Chatbot(label="Workflow Chat") + msg = gr.Textbox(label="Your Input") + submit_btn = gr.Button("Submit") + clear_btn = gr.Button("Clear Chat") + chat_media_name = gr.Textbox(label="Custom Chat Name(optional)") + save_btn = gr.Button("Save Chat to Database") + + def update_workflow_ui(workflow_name): + if not workflow_name: + return {"current_step": 0, "max_steps": 0, "conversation_id": None}, "", [] + selected_workflow = next((wf for wf in workflows if wf['name'] == workflow_name), None) + if selected_workflow: + num_prompts = len(selected_workflow['prompts']) + context = selected_workflow.get('context', '') + first_prompt = selected_workflow['prompts'][0] + initial_chat = [(None, f"{first_prompt}")] + logging.info(f"Initializing workflow: {workflow_name} with {num_prompts} steps") + return {"current_step": 0, "max_steps": num_prompts, "conversation_id": None}, context, initial_chat + else: + logging.error(f"Selected workflow not found: {workflow_name}") + return {"current_step": 0, "max_steps": 0, "conversation_id": None}, "", [] + + def process_workflow_step(message, history, context, workflow_name, api_endpoint, api_key, workflow_state, + save_conv, temp): + logging.info(f"Process workflow step called with message: {message}") + logging.info(f"Current workflow state: {workflow_state}") + try: + selected_workflow = next((wf for wf in workflows if wf['name'] == workflow_name), None) + if not selected_workflow: + logging.error(f"Selected workflow not found: {workflow_name}") + return history, workflow_state, gr.update(interactive=True) + + current_step = workflow_state["current_step"] + max_steps = workflow_state["max_steps"] + + logging.info(f"Current step: {current_step}, Max steps: {max_steps}") + + if current_step >= max_steps: + logging.info("Workflow completed, disabling input") + return history, workflow_state, gr.update(interactive=False) + + prompt = selected_workflow['prompts'][current_step] + full_message = f"{context}\n\nStep {current_step + 1}: {prompt}\nUser: {message}" + + logging.info(f"Calling chat_wrapper with full_message: {full_message[:100]}...") + bot_message, new_history, new_conversation_id = chat_wrapper( + full_message, history, media_content.value, selected_parts.value, + api_endpoint, api_key, "", workflow_state["conversation_id"], + save_conv, temp, "You are a helpful assistant guiding through a workflow." + ) + + logging.info(f"Received bot_message: {bot_message[:100]}...") + + next_step = current_step + 1 + new_workflow_state = { + "current_step": next_step, + "max_steps": max_steps, + "conversation_id": new_conversation_id + } + + if next_step >= max_steps: + logging.info("Workflow completed after this step") + return new_history, new_workflow_state, gr.update(interactive=False) + else: + next_prompt = selected_workflow['prompts'][next_step] + new_history.append((None, f"Step {next_step + 1}: {next_prompt}")) + logging.info(f"Moving to next step: {next_step}") + return new_history, new_workflow_state, gr.update(interactive=True) + except Exception as e: + logging.error(f"Error in process_workflow_step: {str(e)}") + return history, workflow_state, gr.update(interactive=True) + + workflow_selector.change( + update_workflow_ui, + inputs=[workflow_selector], + outputs=[workflow_state, context_input, chatbot] + ) + + submit_btn.click( + process_workflow_step, + inputs=[msg, chatbot, context_input, workflow_selector, api_selector, api_key_input, workflow_state, + save_conversation, temperature], + outputs=[chatbot, workflow_state, msg] + ).then( + lambda: gr.update(value=""), + outputs=[msg] + ) + + clear_btn.click( + lambda: ([], {"current_step": 0, "max_steps": 0, "conversation_id": None}, ""), + outputs=[chatbot, workflow_state, context_input] + ) + + save_btn.click( + save_chat_history_to_db_wrapper, + inputs=[chatbot, conversation_id, media_content, chat_media_name], + outputs=[conversation_id, gr.Textbox(label="Save Status")] + ) + + search_conversations_btn.click( + search_conversations, + inputs=[conversation_search], + outputs=[previous_conversations] + ) + + load_conversations_btn.click( + lambda: ([], {"current_step": 0, "max_steps": 0, "conversation_id": None}, ""), + outputs=[chatbot, workflow_state, context_input] + ).then( + load_conversation, + inputs=[previous_conversations], + outputs=[chatbot, conversation_id] + ) + + return workflow_selector, api_selector, api_key_input, context_input, chatbot, msg, submit_btn, clear_btn, save_btn + +# +# End of script +############################################################################################################ diff --git a/App_Function_Libraries/Gradio_UI/Chat_ui.py b/App_Function_Libraries/Gradio_UI/Chat_ui.py new file mode 100644 index 0000000000000000000000000000000000000000..a8b55e68ac4f20028a858b4f261263fc3b46ce5d --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Chat_ui.py @@ -0,0 +1,1185 @@ +# Chat_ui.py +# Description: Chat interface functions for Gradio +# +# Imports +import html +import json +import logging +import os +import sqlite3 +from datetime import datetime +# +# External Imports +import gradio as gr +# +# Local Imports +from App_Function_Libraries.Chat import chat, save_chat_history, update_chat_content, save_chat_history_to_db_wrapper +from App_Function_Libraries.DB.DB_Manager import add_chat_message, search_chat_conversations, create_chat_conversation, \ + get_chat_messages, update_chat_message, delete_chat_message, load_preset_prompts, db +from App_Function_Libraries.Gradio_UI.Gradio_Shared import update_dropdown, update_user_prompt + + +# +# +######################################################################################################################## +# +# Functions: + + +def show_edit_message(selected): + if selected: + return gr.update(value=selected[0], visible=True), gr.update(value=selected[1], visible=True), gr.update( + visible=True) + return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False) + + +def show_delete_message(selected): + if selected: + return gr.update(value=selected[1], visible=True), gr.update(visible=True) + return gr.update(visible=False), gr.update(visible=False) + + +def debug_output(media_content, selected_parts): + print(f"Debug - Media Content: {media_content}") + print(f"Debug - Selected Parts: {selected_parts}") + return "" + + +def update_selected_parts(use_content, use_summary, use_prompt): + selected_parts = [] + if use_content: + selected_parts.append("content") + if use_summary: + selected_parts.append("summary") + if use_prompt: + selected_parts.append("prompt") + print(f"Debug - Update Selected Parts: {selected_parts}") + return selected_parts + + +# Old update_user_prompt shim for backwards compatibility +def get_system_prompt(preset_name): + # For backwards compatibility + prompts = update_user_prompt(preset_name) + return prompts["system_prompt"] + +def clear_chat(): + """ + Return empty list for chatbot and None for conversation_id + @return: + """ + return gr.update(value=[]), None + + +def clear_chat_single(): + """ + Clears the chatbot and chat history. + + Returns: + list: Empty list for chatbot messages. + list: Empty list for chat history. + """ + return [], [] + +# FIXME - add additional features.... +def chat_wrapper(message, history, media_content, selected_parts, api_endpoint, api_key, custom_prompt, conversation_id, + save_conversation, temperature, system_prompt, max_tokens=None, top_p=None, frequency_penalty=None, + presence_penalty=None, stop_sequence=None): + try: + if save_conversation: + if conversation_id is None: + # Create a new conversation + media_id = media_content.get('id', None) + conversation_name = f"Chat about {media_content.get('title', 'Unknown Media')} - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" + conversation_id = create_chat_conversation(media_id, conversation_name) + + # Add user message to the database + user_message_id = add_chat_message(conversation_id, "user", message) + + # Include the selected parts and custom_prompt only for the first message + if not history and selected_parts: + message_body = "\n".join(selected_parts) + full_message = f"{custom_prompt}\n\n{message}\n\n{message_body}" + elif custom_prompt: + full_message = f"{custom_prompt}\n\n{message}" + else: + full_message = message + + # Generate bot response + bot_message = chat(full_message, history, media_content, selected_parts, api_endpoint, api_key, custom_prompt, + temperature, system_prompt) + + logging.debug(f"Bot message being returned: {bot_message}") + + if save_conversation: + # Add assistant message to the database + add_chat_message(conversation_id, "assistant", bot_message) + + # Update history + new_history = history + [(message, bot_message)] + + return bot_message, new_history, conversation_id + except Exception as e: + logging.error(f"Error in chat wrapper: {str(e)}") + return "An error occurred.", history, conversation_id + +def search_conversations(query): + try: + conversations = search_chat_conversations(query) + if not conversations: + print(f"Debug - Search Conversations - No results found for query: {query}") + return gr.update(choices=[]) + + conversation_options = [ + (f"{c['conversation_name']} (Media: {c['media_title']}, ID: {c['id']})", c['id']) + for c in conversations + ] + print(f"Debug - Search Conversations - Options: {conversation_options}") + return gr.update(choices=conversation_options) + except Exception as e: + print(f"Debug - Search Conversations - Error: {str(e)}") + return gr.update(choices=[]) + + +def load_conversation(conversation_id): + if not conversation_id: + return [], None + + messages = get_chat_messages(conversation_id) + history = [ + (msg['message'], None) if msg['sender'] == 'user' else (None, msg['message']) + for msg in messages + ] + return history, conversation_id + + +def update_message_in_chat(message_id, new_text, history): + update_chat_message(message_id, new_text) + updated_history = [(msg1, msg2) if msg1[1] != message_id and msg2[1] != message_id + else ((new_text, msg1[1]) if msg1[1] == message_id else (new_text, msg2[1])) + for msg1, msg2 in history] + return updated_history + + +def delete_message_from_chat(message_id, history): + delete_chat_message(message_id) + updated_history = [(msg1, msg2) for msg1, msg2 in history if msg1[1] != message_id and msg2[1] != message_id] + return updated_history + + +def regenerate_last_message(history, media_content, selected_parts, api_endpoint, api_key, custom_prompt, temperature, system_prompt): + if not history: + return history, "No messages to regenerate." + + last_entry = history[-1] + last_user_message, last_bot_message = last_entry + + if last_bot_message is None: + return history, "The last message is not from the bot." + + new_history = history[:-1] + + if not last_user_message: + return new_history, "No user message to regenerate the bot response." + + full_message = last_user_message + + bot_message = chat( + full_message, + new_history, + media_content, + selected_parts, + api_endpoint, + api_key, + custom_prompt, + temperature, + system_prompt + ) + + new_history.append((last_user_message, bot_message)) + + return new_history, "Last message regenerated successfully." + +def create_chat_interface(): + custom_css = """ + .chatbot-container .message-wrap .message { + font-size: 14px !important; + } + """ + with gr.TabItem("Remote LLM Chat (Horizontal)", visible=True): + gr.Markdown("# Chat with a designated LLM Endpoint, using your selected item as starting context") + chat_history = gr.State([]) + media_content = gr.State({}) + selected_parts = gr.State([]) + conversation_id = gr.State(None) + + with gr.Row(): + with gr.Column(scale=1): + search_query_input = gr.Textbox(label="Search Query", placeholder="Enter your search query here...") + search_type_input = gr.Radio(choices=["Title", "URL", "Keyword", "Content"], value="Title", + label="Search By") + search_button = gr.Button("Search") + items_output = gr.Dropdown(label="Select Item", choices=[], interactive=True) + item_mapping = gr.State({}) + with gr.Row(): + use_content = gr.Checkbox(label="Use Content") + use_summary = gr.Checkbox(label="Use Summary") + use_prompt = gr.Checkbox(label="Use Prompt") + save_conversation = gr.Checkbox(label="Save Conversation", value=False, visible=True) + with gr.Row(): + temperature = gr.Slider(label="Temperature", minimum=0.00, maximum=1.0, step=0.05, value=0.7) + with gr.Row(): + conversation_search = gr.Textbox(label="Search Conversations") + with gr.Row(): + search_conversations_btn = gr.Button("Search Conversations") + with gr.Row(): + previous_conversations = gr.Dropdown(label="Select Conversation", choices=[], interactive=True) + with gr.Row(): + load_conversations_btn = gr.Button("Load Selected Conversation") + + api_endpoint = gr.Dropdown(label="Select API Endpoint", + choices=["Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", "DeepSeek", + "Mistral", "OpenRouter", + "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM", "ollama", + "HuggingFace"]) + api_key = gr.Textbox(label="API Key (if required)", type="password") + custom_prompt_checkbox = gr.Checkbox(label="Use a Custom Prompt", + value=False, + visible=True) + preset_prompt_checkbox = gr.Checkbox(label="Use a pre-set Prompt", + value=False, + visible=True) + preset_prompt = gr.Dropdown(label="Select Preset Prompt", + choices=load_preset_prompts(), + visible=False) + user_prompt = gr.Textbox(label="Custom Prompt", + placeholder="Enter custom prompt here", + lines=3, + visible=False) + system_prompt_input = gr.Textbox(label="System Prompt", + value="You are a helpful AI assitant", + lines=3, + visible=False) + with gr.Column(scale=2): + chatbot = gr.Chatbot(height=600, elem_classes="chatbot-container") + msg = gr.Textbox(label="Enter your message") + submit = gr.Button("Submit") + regenerate_button = gr.Button("Regenerate Last Message") + clear_chat_button = gr.Button("Clear Chat") + + edit_message_id = gr.Number(label="Message ID to Edit", visible=False) + edit_message_text = gr.Textbox(label="Edit Message", visible=False) + update_message_button = gr.Button("Update Message", visible=False) + + delete_message_id = gr.Number(label="Message ID to Delete", visible=False) + delete_message_button = gr.Button("Delete Message", visible=False) + + chat_media_name = gr.Textbox(label="Custom Chat Name(optional)") + save_chat_history_to_db = gr.Button("Save Chat History to DataBase") + save_chat_history_as_file = gr.Button("Save Chat History as File") + download_file = gr.File(label="Download Chat History") + save_status = gr.Textbox(label="Save Status", interactive=False) + + # Restore original functionality + search_button.click( + fn=update_dropdown, + inputs=[search_query_input, search_type_input], + outputs=[items_output, item_mapping] + ) + + def save_chat_wrapper(history, conversation_id, media_content): + file_path = save_chat_history(history, conversation_id, media_content) + if file_path: + return file_path, f"Chat history saved successfully as {os.path.basename(file_path)}!" + else: + return None, "Error saving chat history. Please check the logs and try again." + + save_chat_history_as_file.click( + save_chat_wrapper, + inputs=[chatbot, conversation_id, media_content], + outputs=[download_file, save_status] + ) + + def update_prompts(preset_name): + prompts = update_user_prompt(preset_name) + return ( + gr.update(value=prompts["user_prompt"], visible=True), + gr.update(value=prompts["system_prompt"], visible=True) + ) + + def clear_chat(): + return [], None # Return empty list for chatbot and None for conversation_id + + clear_chat_button.click( + clear_chat, + outputs=[chatbot, conversation_id] + ) + preset_prompt.change( + update_prompts, + inputs=preset_prompt, + outputs=[user_prompt, system_prompt_input] + ) + custom_prompt_checkbox.change( + fn=lambda x: (gr.update(visible=x), gr.update(visible=x)), + inputs=[custom_prompt_checkbox], + outputs=[user_prompt, system_prompt_input] + ) + preset_prompt_checkbox.change( + fn=lambda x: gr.update(visible=x), + inputs=[preset_prompt_checkbox], + outputs=[preset_prompt] + ) + submit.click( + chat_wrapper, + inputs=[msg, chatbot, media_content, selected_parts, api_endpoint, api_key, user_prompt, conversation_id, + save_conversation, temperature, system_prompt_input], + outputs=[msg, chatbot, conversation_id] + ).then( # Clear the message box after submission + lambda x: gr.update(value=""), + inputs=[chatbot], + outputs=[msg] + ).then( # Clear the user prompt after the first message + lambda: (gr.update(value=""), gr.update(value="")), + outputs=[user_prompt, system_prompt_input] + ) + + items_output.change( + update_chat_content, + inputs=[items_output, use_content, use_summary, use_prompt, item_mapping], + outputs=[media_content, selected_parts] + ) + use_content.change(update_selected_parts, inputs=[use_content, use_summary, use_prompt], + outputs=[selected_parts]) + use_summary.change(update_selected_parts, inputs=[use_content, use_summary, use_prompt], + outputs=[selected_parts]) + use_prompt.change(update_selected_parts, inputs=[use_content, use_summary, use_prompt], + outputs=[selected_parts]) + items_output.change(debug_output, inputs=[media_content, selected_parts], outputs=[]) + + search_conversations_btn.click( + search_conversations, + inputs=[conversation_search], + outputs=[previous_conversations] + ) + + load_conversations_btn.click( + clear_chat, + outputs=[chatbot, chat_history] + ).then( + load_conversation, + inputs=[previous_conversations], + outputs=[chatbot, conversation_id] + ) + + previous_conversations.change( + load_conversation, + inputs=[previous_conversations], + outputs=[chat_history] + ) + + update_message_button.click( + update_message_in_chat, + inputs=[edit_message_id, edit_message_text, chat_history], + outputs=[chatbot] + ) + + delete_message_button.click( + delete_message_from_chat, + inputs=[delete_message_id, chat_history], + outputs=[chatbot] + ) + + save_chat_history_as_file.click( + save_chat_history, + inputs=[chatbot, conversation_id], + outputs=[download_file] + ) + + save_chat_history_to_db.click( + save_chat_history_to_db_wrapper, + inputs=[chatbot, conversation_id, media_content, chat_media_name], + outputs=[conversation_id, gr.Textbox(label="Save Status")] + ) + + regenerate_button.click( + regenerate_last_message, + inputs=[chatbot, media_content, selected_parts, api_endpoint, api_key, user_prompt, temperature, system_prompt_input], + outputs=[chatbot, save_status] + ) + + chatbot.select(show_edit_message, None, [edit_message_text, edit_message_id, update_message_button]) + chatbot.select(show_delete_message, None, [delete_message_id, delete_message_button]) + + +def create_chat_interface_stacked(): + custom_css = """ + .chatbot-container .message-wrap .message { + font-size: 14px !important; + } + """ + with gr.TabItem("Remote LLM Chat - Stacked", visible=True): + gr.Markdown("# Stacked Chat") + chat_history = gr.State([]) + media_content = gr.State({}) + selected_parts = gr.State([]) + conversation_id = gr.State(None) + + with gr.Row(): + with gr.Column(): + search_query_input = gr.Textbox(label="Search Query", placeholder="Enter your search query here...") + search_type_input = gr.Radio(choices=["Title", "URL", "Keyword", "Content"], value="Title", + label="Search By") + search_button = gr.Button("Search") + items_output = gr.Dropdown(label="Select Item", choices=[], interactive=True) + item_mapping = gr.State({}) + with gr.Row(): + use_content = gr.Checkbox(label="Use Content") + use_summary = gr.Checkbox(label="Use Summary") + use_prompt = gr.Checkbox(label="Use Prompt") + save_conversation = gr.Checkbox(label="Save Conversation", value=False, visible=True) + temp = gr.Slider(label="Temperature", minimum=0.00, maximum=1.0, step=0.05, value=0.7) + with gr.Row(): + conversation_search = gr.Textbox(label="Search Conversations") + with gr.Row(): + previous_conversations = gr.Dropdown(label="Select Conversation", choices=[], interactive=True) + with gr.Row(): + search_conversations_btn = gr.Button("Search Conversations") + load_conversations_btn = gr.Button("Load Selected Conversation") + with gr.Column(): + api_endpoint = gr.Dropdown(label="Select API Endpoint", + choices=["Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", "DeepSeek", + "OpenRouter", "Mistral", "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", + "VLLM", "ollama", "HuggingFace"]) + api_key = gr.Textbox(label="API Key (if required)", type="password") + preset_prompt = gr.Dropdown(label="Select Preset Prompt", + choices=load_preset_prompts(), + visible=True) + system_prompt = gr.Textbox(label="System Prompt", + value="You are a helpful AI assistant.", + lines=3, + visible=True) + user_prompt = gr.Textbox(label="Custom User Prompt", + placeholder="Enter custom prompt here", + lines=3, + visible=True) + gr.Markdown("Scroll down for the chat window...") + with gr.Row(): + with gr.Column(scale=1): + chatbot = gr.Chatbot(height=600, elem_classes="chatbot-container") + msg = gr.Textbox(label="Enter your message") + with gr.Row(): + with gr.Column(): + submit = gr.Button("Submit") + regenerate_button = gr.Button("Regenerate Last Message") + clear_chat_button = gr.Button("Clear Chat") + chat_media_name = gr.Textbox(label="Custom Chat Name(optional)", visible=True) + save_chat_history_to_db = gr.Button("Save Chat History to DataBase") + save_chat_history_as_file = gr.Button("Save Chat History as File") + with gr.Column(): + download_file = gr.File(label="Download Chat History") + + # Restore original functionality + search_button.click( + fn=update_dropdown, + inputs=[search_query_input, search_type_input], + outputs=[items_output, item_mapping] + ) + + def update_prompts(preset_name): + prompts = update_user_prompt(preset_name) + return ( + gr.update(value=prompts["user_prompt"], visible=True), + gr.update(value=prompts["system_prompt"], visible=True) + ) + + clear_chat_button.click( + clear_chat, + outputs=[chatbot, conversation_id] + ) + preset_prompt.change( + update_prompts, + inputs=preset_prompt, + outputs=[user_prompt, system_prompt] + ) + + submit.click( + chat_wrapper, + inputs=[msg, chatbot, media_content, selected_parts, api_endpoint, api_key, user_prompt, + conversation_id, save_conversation, temp, system_prompt], + outputs=[msg, chatbot, conversation_id] + ).then( # Clear the message box after submission + lambda x: gr.update(value=""), + inputs=[chatbot], + outputs=[msg] + ).then( # Clear the user prompt after the first message + lambda: gr.update(value=""), + outputs=[user_prompt, system_prompt] + ) + + items_output.change( + update_chat_content, + inputs=[items_output, use_content, use_summary, use_prompt, item_mapping], + outputs=[media_content, selected_parts] + ) + use_content.change(update_selected_parts, inputs=[use_content, use_summary, use_prompt], + outputs=[selected_parts]) + use_summary.change(update_selected_parts, inputs=[use_content, use_summary, use_prompt], + outputs=[selected_parts]) + use_prompt.change(update_selected_parts, inputs=[use_content, use_summary, use_prompt], + outputs=[selected_parts]) + items_output.change(debug_output, inputs=[media_content, selected_parts], outputs=[]) + + search_conversations_btn.click( + search_conversations, + inputs=[conversation_search], + outputs=[previous_conversations] + ) + + load_conversations_btn.click( + clear_chat, + outputs=[chatbot, chat_history] + ).then( + load_conversation, + inputs=[previous_conversations], + outputs=[chatbot, conversation_id] + ) + + previous_conversations.change( + load_conversation, + inputs=[previous_conversations], + outputs=[chat_history] + ) + + save_chat_history_as_file.click( + save_chat_history, + inputs=[chatbot, conversation_id], + outputs=[download_file] + ) + + save_chat_history_to_db.click( + save_chat_history_to_db_wrapper, + inputs=[chatbot, conversation_id, media_content, chat_media_name], + outputs=[conversation_id, gr.Textbox(label="Save Status")] + ) + + regenerate_button.click( + regenerate_last_message, + inputs=[chatbot, media_content, selected_parts, api_endpoint, api_key, user_prompt, temp, system_prompt], + outputs=[chatbot, gr.Textbox(label="Regenerate Status")] + ) + + +# FIXME - System prompts +def create_chat_interface_multi_api(): + custom_css = """ + .chatbot-container .message-wrap .message { + font-size: 14px !important; + } + .chat-window { + height: 400px; + overflow-y: auto; + } + """ + with gr.TabItem("One Prompt - Multiple APIs", visible=True): + gr.Markdown("# One Prompt but Multiple APIs Chat Interface") + + with gr.Row(): + with gr.Column(scale=1): + search_query_input = gr.Textbox(label="Search Query", placeholder="Enter your search query here...") + search_type_input = gr.Radio(choices=["Title", "URL", "Keyword", "Content"], value="Title", + label="Search By") + search_button = gr.Button("Search") + items_output = gr.Dropdown(label="Select Item", choices=[], interactive=True) + item_mapping = gr.State({}) + with gr.Row(): + use_content = gr.Checkbox(label="Use Content") + use_summary = gr.Checkbox(label="Use Summary") + use_prompt = gr.Checkbox(label="Use Prompt") + with gr.Column(): + preset_prompt = gr.Dropdown(label="Select Preset Prompt", choices=load_preset_prompts(), visible=True) + system_prompt = gr.Textbox(label="System Prompt", value="You are a helpful AI assistant.", lines=5) + user_prompt = gr.Textbox(label="Modify Prompt (Prefixed to your message every time)", lines=5, value="", visible=True) + + with gr.Row(): + chatbots = [] + api_endpoints = [] + api_keys = [] + temperatures = [] + regenerate_buttons = [] + for i in range(3): + with gr.Column(): + gr.Markdown(f"### Chat Window {i + 1}") + api_endpoint = gr.Dropdown(label=f"API Endpoint {i + 1}", + choices=["Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", + "DeepSeek", "Mistral", "OpenRouter", "Llama.cpp", "Kobold", + "Ooba", "Tabbyapi", "VLLM", "ollama", "HuggingFace"]) + api_key = gr.Textbox(label=f"API Key {i + 1} (if required)", type="password") + temperature = gr.Slider(label=f"Temperature {i + 1}", minimum=0.0, maximum=1.0, step=0.05, + value=0.7) + chatbot = gr.Chatbot(height=800, elem_classes="chat-window") + regenerate_button = gr.Button(f"Regenerate Last Message {i + 1}") + chatbots.append(chatbot) + api_endpoints.append(api_endpoint) + api_keys.append(api_key) + temperatures.append(temperature) + regenerate_buttons.append(regenerate_button) + + with gr.Row(): + msg = gr.Textbox(label="Enter your message", scale=4) + submit = gr.Button("Submit", scale=1) + clear_chat_button = gr.Button("Clear All Chats") + + # State variables + chat_history = [gr.State([]) for _ in range(3)] + media_content = gr.State({}) + selected_parts = gr.State([]) + conversation_id = gr.State(None) + + # Event handlers + search_button.click( + fn=update_dropdown, + inputs=[search_query_input, search_type_input], + outputs=[items_output, item_mapping] + ) + + preset_prompt.change(update_user_prompt, inputs=preset_prompt, outputs=user_prompt) + + + def clear_all_chats(): + return [[]] * 3 + [[]] * 3 + + clear_chat_button.click( + clear_all_chats, + outputs=chatbots + chat_history + ) + def chat_wrapper_multi(message, custom_prompt, system_prompt, *args): + chat_histories = args[:3] + chatbots = args[3:6] + api_endpoints = args[6:9] + api_keys = args[9:12] + temperatures = args[12:15] + media_content = args[15] + selected_parts = args[16] + + new_chat_histories = [] + new_chatbots = [] + + for i in range(3): + # Call chat_wrapper with dummy values for conversation_id and save_conversation + bot_message, new_history, _ = chat_wrapper( + message, chat_histories[i], media_content, selected_parts, + api_endpoints[i], api_keys[i], custom_prompt, None, # None for conversation_id + False, # False for save_conversation + temperature=temperatures[i], + system_prompt=system_prompt + ) + + new_chatbot = chatbots[i] + [(message, bot_message)] + + new_chat_histories.append(new_history) + new_chatbots.append(new_chatbot) + + return [gr.update(value="")] + new_chatbots + new_chat_histories + + + def regenerate_last_message(chat_history, chatbot, media_content, selected_parts, api_endpoint, api_key, custom_prompt, temperature, system_prompt): + if not chat_history: + return chatbot, chat_history, "No messages to regenerate." + + last_entry = chat_history[-1] + last_user_message, last_bot_message = last_entry + + if last_bot_message is None: + return chatbot, chat_history, "The last message is not from the bot." + + new_history = chat_history[:-1] + + if not last_user_message: + return chatbot[:-1], new_history, "No user message to regenerate the bot response." + + bot_message = chat( + last_user_message, + new_history, + media_content, + selected_parts, + api_endpoint, + api_key, + custom_prompt, + temperature, + system_prompt + ) + + new_history.append((last_user_message, bot_message)) + new_chatbot = chatbot[:-1] + [(last_user_message, bot_message)] + + return new_chatbot, new_history, "Last message regenerated successfully." + + for i in range(3): + regenerate_buttons[i].click( + regenerate_last_message, + inputs=[chat_history[i], chatbots[i], media_content, selected_parts, api_endpoints[i], api_keys[i], user_prompt, temperatures[i], system_prompt], + outputs=[chatbots[i], chat_history[i], gr.Textbox(label=f"Regenerate Status {i + 1}")] + ) + + # In the create_chat_interface_multi_api function: + submit.click( + chat_wrapper_multi, + inputs=[msg, user_prompt, + system_prompt] + chat_history + chatbots + api_endpoints + api_keys + temperatures + + [media_content, selected_parts], + outputs=[msg] + chatbots + chat_history + ).then( + lambda: (gr.update(value=""), gr.update(value="")), + outputs=[msg, user_prompt] + ) + + items_output.change( + update_chat_content, + inputs=[items_output, use_content, use_summary, use_prompt, item_mapping], + outputs=[media_content, selected_parts] + ) + + for checkbox in [use_content, use_summary, use_prompt]: + checkbox.change( + update_selected_parts, + inputs=[use_content, use_summary, use_prompt], + outputs=[selected_parts] + ) + + + +def create_chat_interface_four(): + custom_css = """ + .chatbot-container .message-wrap .message { + font-size: 14px !important; + } + .chat-window { + height: 400px; + overflow-y: auto; + } + """ + + with gr.TabItem("Four Independent API Chats", visible=True): + gr.Markdown("# Four Independent API Chat Interfaces") + + with gr.Row(): + with gr.Column(): + preset_prompt = gr.Dropdown( + label="Select Preset Prompt", + choices=load_preset_prompts(), + visible=True + ) + user_prompt = gr.Textbox( + label="Modify Prompt", + lines=3 + ) + with gr.Column(): + gr.Markdown("Scroll down for the chat windows...") + + chat_interfaces = [] + + def create_single_chat_interface(index, user_prompt_component): + with gr.Column(): + gr.Markdown(f"### Chat Window {index + 1}") + api_endpoint = gr.Dropdown( + label=f"API Endpoint {index + 1}", + choices=[ + "Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", + "DeepSeek", "Mistral", "OpenRouter", "Llama.cpp", "Kobold", + "Ooba", "Tabbyapi", "VLLM", "ollama", "HuggingFace" + ] + ) + api_key = gr.Textbox( + label=f"API Key {index + 1} (if required)", + type="password" + ) + temperature = gr.Slider( + label=f"Temperature {index + 1}", + minimum=0.0, + maximum=1.0, + step=0.05, + value=0.7 + ) + chatbot = gr.Chatbot(height=400, elem_classes="chat-window") + msg = gr.Textbox(label=f"Enter your message for Chat {index + 1}") + submit = gr.Button(f"Submit to Chat {index + 1}") + regenerate_button = gr.Button(f"Regenerate Last Message {index + 1}") + clear_chat_button = gr.Button(f"Clear Chat {index + 1}") + + # State to maintain chat history + chat_history = gr.State([]) + + # Append to chat_interfaces list + chat_interfaces.append({ + 'api_endpoint': api_endpoint, + 'api_key': api_key, + 'temperature': temperature, + 'chatbot': chatbot, + 'msg': msg, + 'submit': submit, + 'regenerate_button': regenerate_button, + 'clear_chat_button': clear_chat_button, + 'chat_history': chat_history + }) + + # Create four chat interfaces arranged in a 2x2 grid + with gr.Row(): + for i in range(2): + with gr.Column(): + for j in range(2): + create_single_chat_interface(i * 2 + j, user_prompt) + + # Update user_prompt based on preset_prompt selection + preset_prompt.change( + fn=update_user_prompt, + inputs=preset_prompt, + outputs=user_prompt + ) + + def chat_wrapper_single(message, chat_history, api_endpoint, api_key, temperature, user_prompt): + logging.debug(f"Chat Wrapper Single - Message: {message}, Chat History: {chat_history}") + + new_msg, new_history, _ = chat_wrapper( + message, + chat_history, + {}, # Empty media_content + [], # Empty selected_parts + api_endpoint, + api_key, + user_prompt, # custom_prompt + None, # conversation_id + False, # save_conversation + temperature, # temperature + system_prompt="", # system_prompt + max_tokens=None, + top_p=None, + frequency_penalty=None, + presence_penalty=None, + stop_sequence=None + ) + if "API request failed" not in new_msg: + chat_history.append((message, new_msg)) + else: + logging.error(f"API request failed: {new_msg}") + + return "", chat_history, chat_history + + def regenerate_last_message(chat_history, api_endpoint, api_key, temperature, user_prompt): + if not chat_history: + return chat_history, chat_history, "No messages to regenerate." + + last_user_message, _ = chat_history[-1] + + new_msg, new_history, _ = chat_wrapper( + last_user_message, + chat_history[:-1], + {}, # Empty media_content + [], # Empty selected_parts + api_endpoint, + api_key, + user_prompt, # custom_prompt + None, # conversation_id + False, # save_conversation + temperature, # temperature + system_prompt="", # system_prompt + max_tokens=None, + top_p=None, + frequency_penalty=None, + presence_penalty=None, + stop_sequence=None + ) + + if "API request failed" not in new_msg: + new_history.append((last_user_message, new_msg)) + return new_history, new_history, "Last message regenerated successfully." + else: + logging.error(f"API request failed during regeneration: {new_msg}") + return chat_history, chat_history, f"Failed to regenerate: {new_msg}" + + # Attach click events for each chat interface + for interface in chat_interfaces: + interface['submit'].click( + chat_wrapper_single, + inputs=[ + interface['msg'], + interface['chat_history'], + interface['api_endpoint'], + interface['api_key'], + interface['temperature'], + user_prompt + ], + outputs=[ + interface['msg'], + interface['chatbot'], + interface['chat_history'] + ] + ) + + interface['regenerate_button'].click( + regenerate_last_message, + inputs=[ + interface['chat_history'], + interface['api_endpoint'], + interface['api_key'], + interface['temperature'], + user_prompt + ], + outputs=[ + interface['chatbot'], + interface['chat_history'], + gr.Textbox(label="Regenerate Status") + ] + ) + + interface['clear_chat_button'].click( + clear_chat_single, + inputs=[], + outputs=[interface['chatbot'], interface['chat_history']] + ) + + +def chat_wrapper_single(message, chat_history, chatbot, api_endpoint, api_key, temperature, media_content, + selected_parts, conversation_id, save_conversation, user_prompt): + new_msg, new_history, new_conv_id = chat_wrapper( + message, chat_history, media_content, selected_parts, + api_endpoint, api_key, user_prompt, conversation_id, + save_conversation, temperature, system_prompt="" + ) + + if new_msg: + updated_chatbot = chatbot + [(message, new_msg)] + else: + updated_chatbot = chatbot + + return new_msg, updated_chatbot, new_history, new_conv_id + + +# FIXME - Finish implementing functions + testing/valdidation +def create_chat_management_tab(): + with gr.TabItem("Chat Management", visible=True): + gr.Markdown("# Chat Management") + + with gr.Row(): + search_query = gr.Textbox(label="Search Conversations") + search_button = gr.Button("Search") + + conversation_list = gr.Dropdown(label="Select Conversation", choices=[]) + conversation_mapping = gr.State({}) + + with gr.Tabs(): + with gr.TabItem("Edit", visible=True): + chat_content = gr.TextArea(label="Chat Content (JSON)", lines=20, max_lines=50) + save_button = gr.Button("Save Changes") + delete_button = gr.Button("Delete Conversation", variant="stop") + + with gr.TabItem("Preview", visible=True): + chat_preview = gr.HTML(label="Chat Preview") + result_message = gr.Markdown("") + + def search_conversations(query): + conversations = search_chat_conversations(query) + choices = [f"{conv['conversation_name']} (Media: {conv['media_title']}, ID: {conv['id']})" for conv in + conversations] + mapping = {choice: conv['id'] for choice, conv in zip(choices, conversations)} + return gr.update(choices=choices), mapping + + def load_conversations(selected, conversation_mapping): + logging.info(f"Selected: {selected}") + logging.info(f"Conversation mapping: {conversation_mapping}") + + try: + if selected and selected in conversation_mapping: + conversation_id = conversation_mapping[selected] + messages = get_chat_messages(conversation_id) + conversation_data = { + "conversation_id": conversation_id, + "messages": messages + } + json_content = json.dumps(conversation_data, indent=2) + + # Create HTML preview + html_preview = "
" + for msg in messages: + sender_style = "background-color: #e6f3ff;" if msg[ + 'sender'] == 'user' else "background-color: #f0f0f0;" + html_preview += f"
" + html_preview += f"{msg['sender']}: {html.escape(msg['message'])}
" + html_preview += f"Timestamp: {msg['timestamp']}" + html_preview += "
" + html_preview += "
" + + logging.info("Returning json_content and html_preview") + return json_content, html_preview + else: + logging.warning("No conversation selected or not in mapping") + return "", "

No conversation selected

" + except Exception as e: + logging.error(f"Error in load_conversations: {str(e)}") + return f"Error: {str(e)}", "

Error loading conversation

" + + def validate_conversation_json(content): + try: + data = json.loads(content) + if not isinstance(data, dict): + return False, "Invalid JSON structure: root should be an object" + if "conversation_id" not in data or not isinstance(data["conversation_id"], int): + return False, "Missing or invalid conversation_id" + if "messages" not in data or not isinstance(data["messages"], list): + return False, "Missing or invalid messages array" + for msg in data["messages"]: + if not all(key in msg for key in ["sender", "message"]): + return False, "Invalid message structure: missing required fields" + return True, data + except json.JSONDecodeError as e: + return False, f"Invalid JSON: {str(e)}" + + def save_conversation(selected, conversation_mapping, content): + if not selected or selected not in conversation_mapping: + return "Please select a conversation before saving.", "

No changes made

" + + conversation_id = conversation_mapping[selected] + is_valid, result = validate_conversation_json(content) + + if not is_valid: + return f"Error: {result}", "

No changes made due to error

" + + conversation_data = result + if conversation_data["conversation_id"] != conversation_id: + return "Error: Conversation ID mismatch.", "

No changes made due to ID mismatch

" + + try: + with db.get_connection() as conn: + conn.execute("BEGIN TRANSACTION") + cursor = conn.cursor() + + # Backup original conversation + cursor.execute("SELECT * FROM ChatMessages WHERE conversation_id = ?", (conversation_id,)) + original_messages = cursor.fetchall() + backup_data = json.dumps({"conversation_id": conversation_id, "messages": original_messages}) + + # You might want to save this backup_data somewhere + + # Delete existing messages + cursor.execute("DELETE FROM ChatMessages WHERE conversation_id = ?", (conversation_id,)) + + # Insert updated messages + for message in conversation_data["messages"]: + cursor.execute(''' + INSERT INTO ChatMessages (conversation_id, sender, message, timestamp) + VALUES (?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP)) + ''', (conversation_id, message["sender"], message["message"], message.get("timestamp"))) + + conn.commit() + + # Create updated HTML preview + html_preview = "
" + for msg in conversation_data["messages"]: + sender_style = "background-color: #e6f3ff;" if msg[ + 'sender'] == 'user' else "background-color: #f0f0f0;" + html_preview += f"
" + html_preview += f"{msg['sender']}: {html.escape(msg['message'])}
" + html_preview += f"Timestamp: {msg.get('timestamp', 'N/A')}" + html_preview += "
" + html_preview += "
" + + return "Conversation updated successfully.", html_preview + except sqlite3.Error as e: + conn.rollback() + logging.error(f"Database error in save_conversation: {e}") + return f"Error updating conversation: {str(e)}", "

Error occurred while saving

" + except Exception as e: + conn.rollback() + logging.error(f"Unexpected error in save_conversation: {e}") + return f"Unexpected error: {str(e)}", "

Unexpected error occurred

" + + def delete_conversation(selected, conversation_mapping): + if not selected or selected not in conversation_mapping: + return "Please select a conversation before deleting.", "

No changes made

", gr.update(choices=[]) + + conversation_id = conversation_mapping[selected] + + try: + with db.get_connection() as conn: + cursor = conn.cursor() + + # Delete messages associated with the conversation + cursor.execute("DELETE FROM ChatMessages WHERE conversation_id = ?", (conversation_id,)) + + # Delete the conversation itself + cursor.execute("DELETE FROM ChatConversations WHERE id = ?", (conversation_id,)) + + conn.commit() + + # Update the conversation list + remaining_conversations = [choice for choice in conversation_mapping.keys() if choice != selected] + updated_mapping = {choice: conversation_mapping[choice] for choice in remaining_conversations} + + return "Conversation deleted successfully.", "

Conversation deleted

", gr.update(choices=remaining_conversations) + except sqlite3.Error as e: + conn.rollback() + logging.error(f"Database error in delete_conversation: {e}") + return f"Error deleting conversation: {str(e)}", "

Error occurred while deleting

", gr.update() + except Exception as e: + conn.rollback() + logging.error(f"Unexpected error in delete_conversation: {e}") + return f"Unexpected error: {str(e)}", "

Unexpected error occurred

", gr.update() + + def parse_formatted_content(formatted_content): + lines = formatted_content.split('\n') + conversation_id = int(lines[0].split(': ')[1]) + timestamp = lines[1].split(': ')[1] + history = [] + current_role = None + current_content = None + for line in lines[3:]: + if line.startswith("Role: "): + if current_role is not None: + history.append({"role": current_role, "content": ["", current_content]}) + current_role = line.split(': ')[1] + elif line.startswith("Content: "): + current_content = line.split(': ', 1)[1] + if current_role is not None: + history.append({"role": current_role, "content": ["", current_content]}) + return json.dumps({ + "conversation_id": conversation_id, + "timestamp": timestamp, + "history": history + }, indent=2) + + search_button.click( + search_conversations, + inputs=[search_query], + outputs=[conversation_list, conversation_mapping] + ) + + conversation_list.change( + load_conversations, + inputs=[conversation_list, conversation_mapping], + outputs=[chat_content, chat_preview] + ) + + save_button.click( + save_conversation, + inputs=[conversation_list, conversation_mapping, chat_content], + outputs=[result_message, chat_preview] + ) + + delete_button.click( + delete_conversation, + inputs=[conversation_list, conversation_mapping], + outputs=[result_message, chat_preview, conversation_list] + ) + + return search_query, search_button, conversation_list, conversation_mapping, chat_content, save_button, delete_button, result_message, chat_preview + + + +# Mock function to simulate LLM processing +def process_with_llm(workflow, context, prompt, api_endpoint, api_key): + api_key_snippet = api_key[:5] + "..." if api_key else "Not provided" + return f"LLM output using {api_endpoint} (API Key: {api_key_snippet}) for {workflow} with context: {context[:30]}... and prompt: {prompt[:30]}..." + + +# +# End of Chat_ui.py +####################################################################################################################### \ No newline at end of file diff --git a/App_Function_Libraries/Gradio_UI/Config_tab.py b/App_Function_Libraries/Gradio_UI/Config_tab.py new file mode 100644 index 0000000000000000000000000000000000000000..99b97367ef2e2e610d0cbe0fead70ffca530369d --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Config_tab.py @@ -0,0 +1,51 @@ +import gradio as gr +import configparser + +# FIXME +CONFIG_PATH = './Config_Files/config.txt' + +def load_config(): + config = configparser.ConfigParser() + config.read(CONFIG_PATH) + return config + +def save_config(config): + with open(CONFIG_PATH, 'w') as configfile: + config.write(configfile) + +def get_config_as_text(): + with open(CONFIG_PATH, 'r') as file: + content = file.read() + return content, "Config refreshed successfully" + +def save_config_from_text(text): + with open(CONFIG_PATH, 'w') as file: + file.write(text) + return "Config saved successfully" + + +def create_config_editor_tab(): + with gr.TabItem("Edit Config", visible=True): + gr.Markdown("# Edit Configuration File") + + with gr.Row(): + with gr.Column(): + refresh_button = gr.Button("Refresh Config") + + with gr.Column(): + config_text = gr.TextArea(label="Full Config", lines=30) + save_text_button = gr.Button("Save Config") + + with gr.Row(): + output = gr.Textbox(label="Output") + + # Event handlers + refresh_button.click(get_config_as_text, inputs=[], outputs=[config_text, output]) + + config_text.change(lambda: None, None, None) # Dummy handler to enable changes + save_text_button.click(save_config_from_text, inputs=[config_text], outputs=[output]) + + # Initialize the interface + config_text.value = get_config_as_text()[0] # Only set the config text, not the output message + + return refresh_button, config_text, save_text_button, output diff --git a/App_Function_Libraries/Gradio_UI/Embeddings_tab.py b/App_Function_Libraries/Gradio_UI/Embeddings_tab.py new file mode 100644 index 0000000000000000000000000000000000000000..3f4841f9c8b52b50bcc643ed7239c123f33dd003 --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Embeddings_tab.py @@ -0,0 +1,508 @@ +# Embeddings_tabc.py +# Description: This file contains the code for the RAG Chat tab in the Gradio UI +# +# Imports +import json +import logging +# +# External Imports +import gradio as gr +import numpy as np +from tqdm import tqdm +# +# Local Imports +from App_Function_Libraries.DB.DB_Manager import get_all_content_from_database +from App_Function_Libraries.RAG.ChromaDB_Library import chroma_client, \ + store_in_chroma, situate_context +from App_Function_Libraries.RAG.Embeddings_Create import create_embedding, create_embeddings_batch +from App_Function_Libraries.Chunk_Lib import improved_chunking_process, chunk_for_embedding +# +######################################################################################################################## +# +# Functions: + +def create_embeddings_tab(): + with gr.TabItem("Create Embeddings", visible=True): + gr.Markdown("# Create Embeddings for All Content") + + with gr.Row(): + with gr.Column(): + embedding_provider = gr.Radio( + choices=["huggingface", "local", "openai"], + label="Select Embedding Provider", + value="huggingface" + ) + gr.Markdown("Note: Local provider requires a running Llama.cpp/llamafile server.") + gr.Markdown("OpenAI provider requires a valid API key.") + + huggingface_model = gr.Dropdown( + choices=[ + "jinaai/jina-embeddings-v3", + "Alibaba-NLP/gte-large-en-v1.5", + "dunzhang/setll_en_400M_v5", + "custom" + ], + label="Hugging Face Model", + value="jinaai/jina-embeddings-v3", + visible=True + ) + + openai_model = gr.Dropdown( + choices=[ + "text-embedding-3-small", + "text-embedding-3-large" + ], + label="OpenAI Embedding Model", + value="text-embedding-3-small", + visible=False + ) + + custom_embedding_model = gr.Textbox( + label="Custom Embedding Model", + placeholder="Enter your custom embedding model name here", + visible=False + ) + + embedding_api_url = gr.Textbox( + label="API URL (for local provider)", + value="http://localhost:8080/embedding", + visible=False + ) + + # Add chunking options + chunking_method = gr.Dropdown( + choices=["words", "sentences", "paragraphs", "tokens", "semantic"], + label="Chunking Method", + value="words" + ) + max_chunk_size = gr.Slider( + minimum=1, maximum=8000, step=1, value=500, + label="Max Chunk Size" + ) + chunk_overlap = gr.Slider( + minimum=0, maximum=4000, step=1, value=200, + label="Chunk Overlap" + ) + adaptive_chunking = gr.Checkbox( + label="Use Adaptive Chunking", + value=False + ) + + create_button = gr.Button("Create Embeddings") + + with gr.Column(): + status_output = gr.Textbox(label="Status", lines=10) + + def update_provider_options(provider): + if provider == "huggingface": + return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False) + elif provider == "local": + return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True) + else: # OpenAI + return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False) + + def update_huggingface_options(model): + if model == "custom": + return gr.update(visible=True) + else: + return gr.update(visible=False) + + embedding_provider.change( + fn=update_provider_options, + inputs=[embedding_provider], + outputs=[huggingface_model, openai_model, custom_embedding_model, embedding_api_url] + ) + + huggingface_model.change( + fn=update_huggingface_options, + inputs=[huggingface_model], + outputs=[custom_embedding_model] + ) + + def create_all_embeddings(provider, hf_model, openai_model, custom_model, api_url, method, max_size, overlap, adaptive): + try: + all_content = get_all_content_from_database() + if not all_content: + return "No content found in the database." + + chunk_options = { + 'method': method, + 'max_size': max_size, + 'overlap': overlap, + 'adaptive': adaptive + } + + collection_name = "all_content_embeddings" + collection = chroma_client.get_or_create_collection(name=collection_name) + + # Determine the model to use + if provider == "huggingface": + model = custom_model if hf_model == "custom" else hf_model + elif provider == "openai": + model = openai_model + else: + model = custom_model + + for item in all_content: + media_id = item['id'] + text = item['content'] + + chunks = improved_chunking_process(text, chunk_options) + for i, chunk in enumerate(chunks): + chunk_text = chunk['text'] + chunk_id = f"doc_{media_id}_chunk_{i}" + + existing = collection.get(ids=[chunk_id]) + if existing['ids']: + continue + + embedding = create_embedding(chunk_text, provider, model, api_url) + metadata = { + "media_id": str(media_id), + "chunk_index": i, + "total_chunks": len(chunks), + "chunking_method": method, + "max_chunk_size": max_size, + "chunk_overlap": overlap, + "adaptive_chunking": adaptive, + "embedding_model": model, + "embedding_provider": provider, + **chunk['metadata'] + } + store_in_chroma(collection_name, [chunk_text], [embedding], [chunk_id], [metadata]) + + return "Embeddings created and stored successfully for all content." + except Exception as e: + logging.error(f"Error during embedding creation: {str(e)}") + return f"Error: {str(e)}" + + create_button.click( + fn=create_all_embeddings, + inputs=[embedding_provider, huggingface_model, openai_model, custom_embedding_model, embedding_api_url, + chunking_method, max_chunk_size, chunk_overlap, adaptive_chunking], + outputs=status_output + ) + + +def create_view_embeddings_tab(): + with gr.TabItem("View/Update Embeddings", visible=True): + gr.Markdown("# View and Update Embeddings") + item_mapping = gr.State({}) + with gr.Row(): + with gr.Column(): + item_dropdown = gr.Dropdown(label="Select Item", choices=[], interactive=True) + refresh_button = gr.Button("Refresh Item List") + embedding_status = gr.Textbox(label="Embedding Status", interactive=False) + embedding_preview = gr.Textbox(label="Embedding Preview", interactive=False, lines=5) + embedding_metadata = gr.Textbox(label="Embedding Metadata", interactive=False, lines=10) + + with gr.Column(): + create_new_embedding_button = gr.Button("Create New Embedding") + embedding_provider = gr.Radio( + choices=["huggingface", "local", "openai"], + label="Select Embedding Provider", + value="huggingface" + ) + gr.Markdown("Note: Local provider requires a running Llama.cpp/llamafile server.") + gr.Markdown("OpenAI provider requires a valid API key.") + + huggingface_model = gr.Dropdown( + choices=[ + "jinaai/jina-embeddings-v3", + "Alibaba-NLP/gte-large-en-v1.5", + "dunzhang/stella_en_400M_v5", + "custom" + ], + label="Hugging Face Model", + value="jinaai/jina-embeddings-v3", + visible=True + ) + + openai_model = gr.Dropdown( + choices=[ + "text-embedding-3-small", + "text-embedding-3-large" + ], + label="OpenAI Embedding Model", + value="text-embedding-3-small", + visible=False + ) + + custom_embedding_model = gr.Textbox( + label="Custom Embedding Model", + placeholder="Enter your custom embedding model name here", + visible=False + ) + + embedding_api_url = gr.Textbox( + label="API URL (for local provider)", + value="http://localhost:8080/embedding", + visible=False + ) + chunking_method = gr.Dropdown( + choices=["words", "sentences", "paragraphs", "tokens", "semantic"], + label="Chunking Method", + value="words" + ) + max_chunk_size = gr.Slider( + minimum=1, maximum=8000, step=5, value=500, + label="Max Chunk Size" + ) + chunk_overlap = gr.Slider( + minimum=0, maximum=5000, step=5, value=200, + label="Chunk Overlap" + ) + adaptive_chunking = gr.Checkbox( + label="Use Adaptive Chunking", + value=False + ) + contextual_api_choice = gr.Dropdown( + choices=["Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", "DeepSeek", "Mistral", "OpenRouter", "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM", "ollama", "HuggingFace"], + label="Select API for Contextualized Embeddings", + value="OpenAI" + ) + use_contextual_embeddings = gr.Checkbox( + label="Use Contextual Embeddings", + value=True + ) + contextual_api_key = gr.Textbox(label="API Key", lines=1) + + def get_items_with_embedding_status(): + try: + items = get_all_content_from_database() + collection = chroma_client.get_or_create_collection(name="all_content_embeddings") + choices = [] + new_item_mapping = {} + for item in items: + try: + result = collection.get(ids=[f"doc_{item['id']}_chunk_0"]) + embedding_exists = result is not None and result.get('ids') and len(result['ids']) > 0 + status = "Embedding exists" if embedding_exists else "No embedding" + except Exception as e: + print(f"Error checking embedding for item {item['id']}: {str(e)}") + status = "Error checking" + choice = f"{item['title']} ({status})" + choices.append(choice) + new_item_mapping[choice] = item['id'] + return gr.update(choices=choices), new_item_mapping + except Exception as e: + print(f"Error in get_items_with_embedding_status: {str(e)}") + return gr.update(choices=["Error: Unable to fetch items"]), {} + + def update_provider_options(provider): + if provider == "huggingface": + return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False) + elif provider == "local": + return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True) + else: # OpenAI + return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False) + + def update_huggingface_options(model): + if model == "custom": + return gr.update(visible=True) + else: + return gr.update(visible=False) + + def check_embedding_status(selected_item, item_mapping): + if not selected_item: + return "Please select an item", "", "" + + try: + item_id = item_mapping.get(selected_item) + if item_id is None: + return f"Invalid item selected: {selected_item}", "", "" + + item_title = selected_item.rsplit(' (', 1)[0] + collection = chroma_client.get_or_create_collection(name="all_content_embeddings") + + result = collection.get(ids=[f"doc_{item_id}_chunk_0"], include=["embeddings", "metadatas"]) + logging.info(f"ChromaDB result for item '{item_title}' (ID: {item_id}): {result}") + + if not result['ids']: + return f"No embedding found for item '{item_title}' (ID: {item_id})", "", "" + + if not result['embeddings'] or not result['embeddings'][0]: + return f"Embedding data missing for item '{item_title}' (ID: {item_id})", "", "" + + embedding = result['embeddings'][0] + metadata = result['metadatas'][0] if result['metadatas'] else {} + embedding_preview = str(embedding[:50]) + status = f"Embedding exists for item '{item_title}' (ID: {item_id})" + return status, f"First 50 elements of embedding:\n{embedding_preview}", json.dumps(metadata, indent=2) + + except Exception as e: + logging.error(f"Error in check_embedding_status: {str(e)}") + return f"Error processing item: {selected_item}. Details: {str(e)}", "", "" + + def create_new_embedding_for_item(selected_item, provider, hf_model, openai_model, custom_model, api_url, + method, max_size, overlap, adaptive, + item_mapping, use_contextual, contextual_api_choice=None): + if not selected_item: + return "Please select an item", "", "" + + try: + item_id = item_mapping.get(selected_item) + if item_id is None: + return f"Invalid item selected: {selected_item}", "", "" + + items = get_all_content_from_database() + item = next((item for item in items if item['id'] == item_id), None) + if not item: + return f"Item not found: {item_id}", "", "" + + chunk_options = { + 'method': method, + 'max_size': max_size, + 'overlap': overlap, + 'adaptive': adaptive + } + + logging.info(f"Chunking content for item: {item['title']} (ID: {item_id})") + chunks = chunk_for_embedding(item['content'], item['title'], chunk_options) + collection_name = "all_content_embeddings" + collection = chroma_client.get_or_create_collection(name=collection_name) + + # Delete existing embeddings for this item + existing_ids = [f"doc_{item_id}_chunk_{i}" for i in range(len(chunks))] + collection.delete(ids=existing_ids) + logging.info(f"Deleted {len(existing_ids)} existing embeddings for item {item_id}") + + texts, ids, metadatas = [], [], [] + chunk_count = 0 + logging.info("Generating contextual summaries and preparing chunks for embedding") + for i, chunk in enumerate(chunks): + chunk_text = chunk['text'] + chunk_metadata = chunk['metadata'] + if use_contextual: + logging.debug(f"Generating contextual summary for chunk {chunk_count}") + context = situate_context(contextual_api_choice, item['content'], chunk_text) + contextualized_text = f"{chunk_text}\n\nContextual Summary: {context}" + else: + contextualized_text = chunk_text + context = None + + chunk_id = f"doc_{item_id}_chunk_{i}" + + # Determine the model to use + if provider == "huggingface": + model = custom_model if hf_model == "custom" else hf_model + elif provider == "openai": + model = openai_model + else: + model = custom_model + + metadata = { + "media_id": str(item_id), + "chunk_index": i, + "total_chunks": len(chunks), + "chunking_method": method, + "max_chunk_size": max_size, + "chunk_overlap": overlap, + "adaptive_chunking": adaptive, + "embedding_model": model, + "embedding_provider": provider, + "original_text": chunk_text, + "use_contextual_embeddings": use_contextual, + "contextual_summary": context, + **chunk_metadata + } + + texts.append(contextualized_text) + ids.append(chunk_id) + metadatas.append(metadata) + chunk_count += 1 + + # Create embeddings in batch + logging.info(f"Creating embeddings for {len(texts)} chunks") + embeddings = create_embeddings_batch(texts, provider, model, api_url) + + # Store in Chroma + store_in_chroma(collection_name, texts, embeddings, ids, metadatas) + + # Create a preview of the first embedding + if isinstance(embeddings, np.ndarray) and embeddings.size > 0: + embedding_preview = str(embeddings[0][:50]) + elif isinstance(embeddings, list) and len(embeddings) > 0: + embedding_preview = str(embeddings[0][:50]) + else: + embedding_preview = "No embeddings created" + + # Return status message + status = f"New embeddings created and stored for item: {item['title']} (ID: {item_id})" + + # Add contextual summaries to status message if enabled + if use_contextual: + status += " (with contextual summaries)" + + # Return status message, embedding preview, and metadata + return status, f"First 50 elements of new embedding:\n{embedding_preview}", json.dumps(metadatas[0], + indent=2) + except Exception as e: + logging.error(f"Error in create_new_embedding_for_item: {str(e)}", exc_info=True) + return f"Error creating embedding: {str(e)}", "", "" + + refresh_button.click( + get_items_with_embedding_status, + outputs=[item_dropdown, item_mapping] + ) + item_dropdown.change( + check_embedding_status, + inputs=[item_dropdown, item_mapping], + outputs=[embedding_status, embedding_preview, embedding_metadata] + ) + create_new_embedding_button.click( + create_new_embedding_for_item, + inputs=[item_dropdown, embedding_provider, huggingface_model, openai_model, custom_embedding_model, embedding_api_url, + chunking_method, max_chunk_size, chunk_overlap, adaptive_chunking, item_mapping, + use_contextual_embeddings, contextual_api_choice], + outputs=[embedding_status, embedding_preview, embedding_metadata] + ) + embedding_provider.change( + update_provider_options, + inputs=[embedding_provider], + outputs=[huggingface_model, openai_model, custom_embedding_model, embedding_api_url] + ) + huggingface_model.change( + update_huggingface_options, + inputs=[huggingface_model], + outputs=[custom_embedding_model] + ) + + return (item_dropdown, refresh_button, embedding_status, embedding_preview, embedding_metadata, + create_new_embedding_button, embedding_provider, huggingface_model, openai_model, custom_embedding_model, embedding_api_url, + chunking_method, max_chunk_size, chunk_overlap, adaptive_chunking, + use_contextual_embeddings, contextual_api_choice, contextual_api_key) + + +def create_purge_embeddings_tab(): + with gr.TabItem("Purge Embeddings", visible=True): + gr.Markdown("# Purge Embeddings") + + with gr.Row(): + with gr.Column(): + purge_button = gr.Button("Purge All Embeddings") + with gr.Column(): + status_output = gr.Textbox(label="Status", lines=10) + + def purge_all_embeddings(): + try: + # It came to me in a dream....I literally don't remember how the fuck this works, cant find documentation... + collection_name = "all_content_embeddings" + chroma_client.delete_collection(collection_name) + chroma_client.create_collection(collection_name) + logging.info(f"All embeddings have been purged successfully.") + return "All embeddings have been purged successfully." + except Exception as e: + logging.error(f"Error during embedding purge: {str(e)}") + return f"Error: {str(e)}" + + purge_button.click( + fn=purge_all_embeddings, + outputs=status_output + ) + + + +# +# End of file +######################################################################################################################## diff --git a/App_Function_Libraries/Gradio_UI/Evaluations_Benchmarks_tab.py b/App_Function_Libraries/Gradio_UI/Evaluations_Benchmarks_tab.py new file mode 100644 index 0000000000000000000000000000000000000000..f1ffbc69ecbcf8786493397cf6ed45931e561f13 --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Evaluations_Benchmarks_tab.py @@ -0,0 +1,60 @@ +################################################################################################### +# Evaluations_Benchmarks_tab.py - Gradio code for G-Eval testing +# We will use the G-Eval API to evaluate the quality of the generated summaries. + +import gradio as gr +from App_Function_Libraries.Benchmarks_Evaluations.ms_g_eval import run_geval + +def create_geval_tab(): + with gr.Tab("G-Eval", visible=True): + gr.Markdown("# G-Eval Summarization Evaluation") + with gr.Row(): + with gr.Column(): + document_input = gr.Textbox(label="Source Document", lines=10) + summary_input = gr.Textbox(label="Summary", lines=5) + api_name_input = gr.Dropdown( + choices=["OpenAI", "Anthropic", "Cohere", "Groq", "OpenRouter", "DeepSeek", "HuggingFace", "Mistral", "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM", "Local-LLM", "Ollama"], + label="Select API" + ) + api_key_input = gr.Textbox(label="API Key (if required)", type="password") + evaluate_button = gr.Button("Evaluate Summary") + with gr.Column(): + output = gr.Textbox(label="Evaluation Results", lines=10) + + evaluate_button.click( + fn=run_geval, + inputs=[document_input, summary_input, api_name_input, api_key_input], + outputs=output + ) + + return document_input, summary_input, api_name_input, api_key_input, evaluate_button, output + + +def create_infinite_bench_tab(): + with gr.Tab("Infinite Bench", visible=True): + gr.Markdown("# Infinite Bench Evaluation (Coming Soon)") + with gr.Row(): + with gr.Column(): + api_name_input = gr.Dropdown( + choices=["OpenAI", "Anthropic", "Cohere", "Groq", "OpenRouter", "DeepSeek", "HuggingFace", "Mistral", "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM", "Local-LLM", "Ollama"], + label="Select API" + ) + api_key_input = gr.Textbox(label="API Key (if required)", type="password") + evaluate_button = gr.Button("Evaluate Summary") + with gr.Column(): + output = gr.Textbox(label="Evaluation Results", lines=10) + + # evaluate_button.click( + # fn=run_geval, + # inputs=[api_name_input, api_key_input], + # outputs=output + # ) + + return api_name_input, api_key_input, evaluate_button, output + + +# If you want to run this as a standalone Gradio app +if __name__ == "__main__": + with gr.Blocks() as demo: + create_geval_tab() + demo.launch() diff --git a/App_Function_Libraries/Gradio_UI/Explain_summarize_tab.py b/App_Function_Libraries/Gradio_UI/Explain_summarize_tab.py new file mode 100644 index 0000000000000000000000000000000000000000..37349d8df88886a7f67f47fbbbb175cb76893698 --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Explain_summarize_tab.py @@ -0,0 +1,313 @@ +# Explain_summarize_tab.py +# Gradio UI for explaining and summarizing text +# +# Imports +import logging +# +# External Imports +import gradio as gr + +from App_Function_Libraries.DB.DB_Manager import load_preset_prompts +from App_Function_Libraries.Gradio_UI.Gradio_Shared import update_user_prompt +# +# Local Imports +from App_Function_Libraries.Summarization.Local_Summarization_Lib import summarize_with_llama, summarize_with_kobold, \ + summarize_with_oobabooga, summarize_with_tabbyapi, summarize_with_vllm, summarize_with_local_llm, \ + summarize_with_ollama +from App_Function_Libraries.Summarization.Summarization_General_Lib import summarize_with_openai, summarize_with_anthropic, \ + summarize_with_cohere, summarize_with_groq, summarize_with_openrouter, summarize_with_deepseek, \ + summarize_with_huggingface +# +# +############################################################################################################ +# +# Functions: + +def create_summarize_explain_tab(): + with gr.TabItem("Analyze Text", visible=True): + gr.Markdown("# Analyze / Explain / Summarize Text without ingesting it into the DB") + with gr.Row(): + with gr.Column(): + with gr.Row(): + text_to_work_input = gr.Textbox(label="Text to be Explained or Summarized", + placeholder="Enter the text you want explained or summarized here", + lines=20) + with gr.Row(): + explanation_checkbox = gr.Checkbox(label="Explain Text", value=True) + summarization_checkbox = gr.Checkbox(label="Summarize Text", value=True) + custom_prompt_checkbox = gr.Checkbox(label="Use a Custom Prompt", + value=False, + visible=True) + preset_prompt_checkbox = gr.Checkbox(label="Use a pre-set Prompt", + value=False, + visible=True) + with gr.Row(): + preset_prompt = gr.Dropdown(label="Select Preset Prompt", + choices=load_preset_prompts(), + visible=False) + with gr.Row(): + custom_prompt_input = gr.Textbox(label="Custom Prompt", + placeholder="Enter custom prompt here", + lines=3, + visible=False) + with gr.Row(): + system_prompt_input = gr.Textbox(label="System Prompt", + value="""You are a bulleted notes specialist. [INST]```When creating comprehensive bulleted notes, you should follow these guidelines: Use multiple headings based on the referenced topics, not categories like quotes or terms. Headings should be surrounded by bold formatting and not be listed as bullet points themselves. Leave no space between headings and their corresponding list items underneath. Important terms within the content should be emphasized by setting them in bold font. Any text that ends with a colon should also be bolded. Before submitting your response, review the instructions, and make any corrections necessary to adhered to the specified format. Do not reference these instructions within the notes.``` \nBased on the content between backticks create comprehensive bulleted notes.[/INST] + **Bulleted Note Creation Guidelines** + + **Headings**: + - Based on referenced topics, not categories like quotes or terms + - Surrounded by **bold** formatting + - Not listed as bullet points + - No space between headings and list items underneath + + **Emphasis**: + - **Important terms** set in bold font + - **Text ending in a colon**: also bolded + + **Review**: + - Ensure adherence to specified format + - Do not reference these instructions in your response.[INST] {{ .Prompt }} [/INST] + """, + lines=3, + visible=False, + interactive=True) + api_endpoint = gr.Dropdown( + choices=[None, "Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", "DeepSeek", "Mistral", + "OpenRouter", + "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM", "ollama", "HuggingFace", "Custom-OpenAI-API"], + value=None, + label="API to be used for request (Mandatory)" + ) + with gr.Row(): + api_key_input = gr.Textbox(label="API Key (if required)", placeholder="Enter your API key here", + type="password") + with gr.Row(): + explain_summarize_button = gr.Button("Explain/Summarize") + + with gr.Column(): + summarization_output = gr.Textbox(label="Summary:", lines=20) + explanation_output = gr.Textbox(label="Explanation:", lines=20) + custom_prompt_output = gr.Textbox(label="Custom Prompt:", lines=20, visible=True) + + custom_prompt_checkbox.change( + fn=lambda x: (gr.update(visible=x), gr.update(visible=x)), + inputs=[custom_prompt_checkbox], + outputs=[custom_prompt_input, system_prompt_input] + ) + preset_prompt_checkbox.change( + fn=lambda x: gr.update(visible=x), + inputs=[preset_prompt_checkbox], + outputs=[preset_prompt] + ) + + def update_prompts(preset_name): + prompts = update_user_prompt(preset_name) + return ( + gr.update(value=prompts["user_prompt"], visible=True), + gr.update(value=prompts["system_prompt"], visible=True) + ) + + preset_prompt.change( + update_prompts, + inputs=preset_prompt, + outputs=[custom_prompt_input, system_prompt_input] + ) + + explain_summarize_button.click( + fn=summarize_explain_text, + inputs=[text_to_work_input, api_endpoint, api_key_input, summarization_checkbox, explanation_checkbox, custom_prompt_input, system_prompt_input], + outputs=[summarization_output, explanation_output, custom_prompt_output] + ) + + +def summarize_explain_text(message, api_endpoint, api_key, summarization, explanation, custom_prompt, custom_system_prompt,): + global custom_prompt_output + summarization_response = None + explanation_response = None + temp = 0.7 + try: + logging.info(f"Debug - summarize_explain_text Function - Message: {message}") + logging.info(f"Debug - summarize_explain_text Function - API Endpoint: {api_endpoint}") + + # Prepare the input for the API + input_data = f"User: {message}\n" + # Print first 500 chars + logging.info(f"Debug - Chat Function - Input Data: {input_data[:500]}...") + logging.debug(f"Debug - Chat Function - API Key: {api_key[:10]}") + user_prompt = " " + if not api_endpoint: + return "Please select an API endpoint", "Please select an API endpoint" + try: + if summarization: + system_prompt = """You are a bulleted notes specialist. [INST]```When creating comprehensive bulleted notes, you should follow these guidelines: Use multiple headings based on the referenced topics, not categories like quotes or terms. Headings should be surrounded by bold formatting and not be listed as bullet points themselves. Leave no space between headings and their corresponding list items underneath. Important terms within the content should be emphasized by setting them in bold font. Any text that ends with a colon should also be bolded. Before submitting your response, review the instructions, and make any corrections necessary to adhered to the specified format. Do not reference these instructions within the notes.``` \nBased on the content between backticks create comprehensive bulleted notes.[/INST] + **Bulleted Note Creation Guidelines** + + **Headings**: + - Based on referenced topics, not categories like quotes or terms + - Surrounded by **bold** formatting + - Not listed as bullet points + - No space between headings and list items underneath + + **Emphasis**: + - **Important terms** set in bold font + - **Text ending in a colon**: also bolded + + **Review**: + - Ensure adherence to specified format + - Do not reference these instructions in your response.[INST] {{ .Prompt }} [/INST]""" + + # Use the existing API request code based on the selected endpoint + logging.info(f"Debug - Chat Function - API Endpoint: {api_endpoint}") + if api_endpoint.lower() == 'openai': + summarization_response = summarize_with_openai(api_key, input_data, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "anthropic": + summarization_response = summarize_with_anthropic(api_key, input_data, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "cohere": + summarization_response = summarize_with_cohere(api_key, input_data, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "groq": + summarization_response = summarize_with_groq(api_key, input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "openrouter": + summarization_response = summarize_with_openrouter(api_key, input_data, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "deepseek": + summarization_response = summarize_with_deepseek(api_key, input_data, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "llama.cpp": + summarization_response = summarize_with_llama(input_data, user_prompt, api_key, temp, system_prompt) + elif api_endpoint.lower() == "kobold": + summarization_response = summarize_with_kobold(input_data, api_key, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "ooba": + summarization_response = summarize_with_oobabooga(input_data, api_key, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "tabbyapi": + summarization_response = summarize_with_tabbyapi(input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "vllm": + summarization_response = summarize_with_vllm(input_data, user_prompt, system_prompt) + elif api_endpoint.lower() == "local-llm": + summarization_response = summarize_with_local_llm(input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "huggingface": + summarization_response = summarize_with_huggingface(api_key, input_data, user_prompt, + temp) # , system_prompt) + elif api_endpoint.lower() == "ollama": + summarization_response = summarize_with_ollama(input_data, user_prompt, None, api_key, temp, system_prompt) + else: + raise ValueError(f"Unsupported API endpoint: {api_endpoint}") + except Exception as e: + logging.error(f"Error in summarization: {str(e)}") + response1 = f"An error occurred during summarization: {str(e)}" + + try: + if explanation: + system_prompt = """You are a professional teacher. Please explain the content presented in an easy to digest fashion so that a non-specialist may understand it.""" + # Use the existing API request code based on the selected endpoint + logging.info(f"Debug - Chat Function - API Endpoint: {api_endpoint}") + if api_endpoint.lower() == 'openai': + explanation_response = summarize_with_openai(api_key, input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "anthropic": + explanation_response = summarize_with_anthropic(api_key, input_data, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "cohere": + explanation_response = summarize_with_cohere(api_key, input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "groq": + explanation_response = summarize_with_groq(api_key, input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "openrouter": + explanation_response = summarize_with_openrouter(api_key, input_data, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "deepseek": + explanation_response = summarize_with_deepseek(api_key, input_data, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "llama.cpp": + explanation_response = summarize_with_llama(input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "kobold": + explanation_response = summarize_with_kobold(input_data, api_key, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "ooba": + explanation_response = summarize_with_oobabooga(input_data, api_key, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "tabbyapi": + explanation_response = summarize_with_tabbyapi(input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "vllm": + explanation_response = summarize_with_vllm(input_data, user_prompt, system_prompt) + elif api_endpoint.lower() == "local-llm": + explanation_response = summarize_with_local_llm(input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "huggingface": + explanation_response = summarize_with_huggingface(api_key, input_data, user_prompt, + temp) # , system_prompt) + elif api_endpoint.lower() == "ollama": + explanation_response = summarize_with_ollama(input_data, user_prompt, temp, system_prompt) + else: + raise ValueError(f"Unsupported API endpoint: {api_endpoint}") + except Exception as e: + logging.error(f"Error in summarization: {str(e)}") + response2 = f"An error occurred during summarization: {str(e)}" + + try: + if custom_prompt: + system_prompt = custom_system_prompt + user_prompt = custom_prompt + input_data + # Use the existing API request code based on the selected endpoint + logging.info(f"Debug - Chat Function - API Endpoint: {api_endpoint}") + if api_endpoint.lower() == 'openai': + custom_prompt_output = summarize_with_openai(api_key, input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "anthropic": + custom_prompt_output = summarize_with_anthropic(api_key, input_data, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "cohere": + custom_prompt_output = summarize_with_cohere(api_key, input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "groq": + custom_prompt_output = summarize_with_groq(api_key, input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "openrouter": + custom_prompt_output = summarize_with_openrouter(api_key, input_data, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "deepseek": + custom_prompt_output = summarize_with_deepseek(api_key, input_data, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "llama.cpp": + custom_prompt_output = summarize_with_llama(input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "kobold": + custom_prompt_output = summarize_with_kobold(input_data, api_key, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "ooba": + custom_prompt_output = summarize_with_oobabooga(input_data, api_key, user_prompt, temp, + system_prompt) + elif api_endpoint.lower() == "tabbyapi": + custom_prompt_output = summarize_with_tabbyapi(input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "vllm": + custom_prompt_output = summarize_with_vllm(input_data, user_prompt, system_prompt) + elif api_endpoint.lower() == "local-llm": + custom_prompt_output = summarize_with_local_llm(input_data, user_prompt, temp, system_prompt) + elif api_endpoint.lower() == "huggingface": + custom_prompt_output = summarize_with_huggingface(api_key, input_data, user_prompt, + temp) # , system_prompt) + elif api_endpoint.lower() == "ollama": + custom_prompt_output = summarize_with_ollama(input_data, user_prompt, temp, system_prompt) + else: + raise ValueError(f"Unsupported API endpoint: {api_endpoint}") + except Exception as e: + logging.error(f"Error in summarization: {str(e)}") + response2 = f"An error occurred during summarization: {str(e)}" + + + if summarization_response: + response1 = f"Summary: {summarization_response}" + else: + response1 = "Summary: No summary requested" + + if explanation_response: + response2 = f"Explanation: {explanation_response}" + else: + response2 = "Explanation: No explanation requested" + + if custom_prompt_output: + response3 = f"Custom Prompt: {custom_prompt_output}" + else: + response3 = "Custom Prompt: No custom prompt requested" + + return response1, response2, response3 + + except Exception as e: + logging.error(f"Error in chat function: {str(e)}") + return f"An error occurred: {str(e)}" \ No newline at end of file diff --git a/App_Function_Libraries/Gradio_UI/Export_Functionality.py b/App_Function_Libraries/Gradio_UI/Export_Functionality.py new file mode 100644 index 0000000000000000000000000000000000000000..2feed8605a614624f6b6246e6379dd7582e15240 --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Export_Functionality.py @@ -0,0 +1,266 @@ +# Export_Functionality.py +# Functionality for exporting items as markdown files +import os +import json +import math +import logging +import shutil +import tempfile +from typing import List, Dict, Optional, Tuple +import gradio as gr +from App_Function_Libraries.DB.DB_Manager import DatabaseError +from App_Function_Libraries.Gradio_UI.Gradio_Shared import fetch_item_details, fetch_items_by_keyword, browse_items + +logger = logging.getLogger(__name__) + +def export_item_as_markdown(media_id: int) -> Tuple[Optional[str], str]: + try: + content, prompt, summary = fetch_item_details(media_id) + title = f"Item {media_id}" # You might want to fetch the actual title + markdown_content = f"# {title}\n\n## Prompt\n{prompt}\n\n## Summary\n{summary}\n\n## Content\n{content}" + + filename = f"export_item_{media_id}.md" + with open(filename, "w", encoding='utf-8') as f: + f.write(markdown_content) + + logger.info(f"Successfully exported item {media_id} to {filename}") + return filename, f"Successfully exported item {media_id} to {filename}" + except Exception as e: + error_message = f"Error exporting item {media_id}: {str(e)}" + logger.error(error_message) + return None, error_message + + +def export_items_by_keyword(keyword: str) -> str: + try: + items = fetch_items_by_keyword(keyword) + if not items: + logger.warning(f"No items found for keyword: {keyword}") + return None + + # Create a temporary directory to store individual markdown files + with tempfile.TemporaryDirectory() as temp_dir: + folder_name = f"export_keyword_{keyword}" + export_folder = os.path.join(temp_dir, folder_name) + os.makedirs(export_folder) + + for item in items: + content, prompt, summary = fetch_item_details(item['id']) + markdown_content = f"# {item['title']}\n\n## Prompt\n{prompt}\n\n## Summary\n{summary}\n\n## Content\n{content}" + + # Create individual markdown file for each item + file_name = f"{item['id']}_{item['title'][:50]}.md" # Limit filename length + file_path = os.path.join(export_folder, file_name) + with open(file_path, "w", encoding='utf-8') as f: + f.write(markdown_content) + + # Create a zip file containing all markdown files + zip_filename = f"{folder_name}.zip" + shutil.make_archive(os.path.join(temp_dir, folder_name), 'zip', export_folder) + + # Move the zip file to a location accessible by Gradio + final_zip_path = os.path.join(os.getcwd(), zip_filename) + shutil.move(os.path.join(temp_dir, zip_filename), final_zip_path) + + logger.info(f"Successfully exported {len(items)} items for keyword '{keyword}' to {zip_filename}") + return final_zip_path + except Exception as e: + logger.error(f"Error exporting items for keyword '{keyword}': {str(e)}") + return None + + +def export_selected_items(selected_items: List[Dict]) -> Tuple[Optional[str], str]: + try: + logger.debug(f"Received selected_items: {selected_items}") + if not selected_items: + logger.warning("No items selected for export") + return None, "No items selected for export" + + markdown_content = "# Selected Items\n\n" + for item in selected_items: + logger.debug(f"Processing item: {item}") + try: + # Check if 'value' is a string (JSON) or already a dictionary + if isinstance(item, str): + item_data = json.loads(item) + elif isinstance(item, dict) and 'value' in item: + item_data = item['value'] if isinstance(item['value'], dict) else json.loads(item['value']) + else: + item_data = item + + logger.debug(f"Item data after processing: {item_data}") + + if 'id' not in item_data: + logger.error(f"'id' not found in item data: {item_data}") + continue + + content, prompt, summary = fetch_item_details(item_data['id']) + markdown_content += f"## {item_data.get('title', 'Item {}'.format(item_data['id']))}\n\n### Prompt\n{prompt}\n\n### Summary\n{summary}\n\n### Content\n{content}\n\n---\n\n" + except Exception as e: + logger.error(f"Error processing item {item}: {str(e)}") + markdown_content += f"## Error\n\nUnable to process this item.\n\n---\n\n" + + filename = "export_selected_items.md" + with open(filename, "w", encoding='utf-8') as f: + f.write(markdown_content) + + logger.info(f"Successfully exported {len(selected_items)} selected items to {filename}") + return filename, f"Successfully exported {len(selected_items)} items to {filename}" + except Exception as e: + error_message = f"Error exporting selected items: {str(e)}" + logger.error(error_message) + return None, error_message + + +def display_search_results_export_tab(search_query: str, search_type: str, page: int = 1, items_per_page: int = 10): + logger.info(f"Searching with query: '{search_query}', type: '{search_type}', page: {page}") + try: + results = browse_items(search_query, search_type) + logger.info(f"browse_items returned {len(results)} results") + + if not results: + return [], f"No results found for query: '{search_query}'", 1, 1 + + total_pages = math.ceil(len(results) / items_per_page) + start_index = (page - 1) * items_per_page + end_index = start_index + items_per_page + paginated_results = results[start_index:end_index] + + checkbox_data = [ + { + "name": f"Name: {item[1]}\nURL: {item[2]}", + "value": {"id": item[0], "title": item[1], "url": item[2]} + } + for item in paginated_results + ] + + logger.info(f"Returning {len(checkbox_data)} items for checkbox (page {page} of {total_pages})") + return checkbox_data, f"Found {len(results)} results (showing page {page} of {total_pages})", page, total_pages + + except DatabaseError as e: + error_message = f"Error in display_search_results_export_tab: {str(e)}" + logger.error(error_message) + return [], error_message, 1, 1 + except Exception as e: + error_message = f"Unexpected error in display_search_results_export_tab: {str(e)}" + logger.error(error_message) + return [], error_message, 1, 1 + + +def create_export_tab(): + with gr.Tab("Search and Export"): + with gr.Row(): + with gr.Column(): + gr.Markdown("# Search and Export Items") + gr.Markdown("Search for items and export them as markdown files") + gr.Markdown("You can also export items by keyword") + search_query = gr.Textbox(label="Search Query") + search_type = gr.Radio(["Title", "URL", "Keyword", "Content"], label="Search By") + search_button = gr.Button("Search") + + with gr.Column(): + prev_button = gr.Button("Previous Page") + next_button = gr.Button("Next Page") + + current_page = gr.State(1) + total_pages = gr.State(1) + + search_results = gr.CheckboxGroup(label="Search Results", choices=[]) + export_selected_button = gr.Button("Export Selected Items") + + keyword_input = gr.Textbox(label="Enter keyword for export") + export_by_keyword_button = gr.Button("Export items by keyword") + + export_output = gr.File(label="Download Exported File") + error_output = gr.Textbox(label="Status/Error Messages", interactive=False) + + def search_and_update(query, search_type, page): + results, message, current, total = display_search_results_export_tab(query, search_type, page) + logger.debug(f"search_and_update results: {results}") + return results, message, current, total, gr.update(choices=results) + + search_button.click( + fn=search_and_update, + inputs=[search_query, search_type, current_page], + outputs=[search_results, error_output, current_page, total_pages, search_results], + show_progress="full" + ) + + + def update_page(current, total, direction): + new_page = max(1, min(total, current + direction)) + return new_page + + prev_button.click( + fn=update_page, + inputs=[current_page, total_pages, gr.State(-1)], + outputs=[current_page] + ).then( + fn=search_and_update, + inputs=[search_query, search_type, current_page], + outputs=[search_results, error_output, current_page, total_pages], + show_progress=True + ) + + next_button.click( + fn=update_page, + inputs=[current_page, total_pages, gr.State(1)], + outputs=[current_page] + ).then( + fn=search_and_update, + inputs=[search_query, search_type, current_page], + outputs=[search_results, error_output, current_page, total_pages], + show_progress=True + ) + + def handle_export_selected(selected_items): + logger.debug(f"Exporting selected items: {selected_items}") + return export_selected_items(selected_items) + + export_selected_button.click( + fn=handle_export_selected, + inputs=[search_results], + outputs=[export_output, error_output], + show_progress="full" + ) + + export_by_keyword_button.click( + fn=export_items_by_keyword, + inputs=[keyword_input], + outputs=[export_output, error_output], + show_progress="full" + ) + + def handle_item_selection(selected_items): + logger.debug(f"Selected items: {selected_items}") + if not selected_items: + return None, "No item selected" + + try: + # Assuming selected_items is a list of dictionaries + selected_item = selected_items[0] + logger.debug(f"First selected item: {selected_item}") + + # Check if 'value' is a string (JSON) or already a dictionary + if isinstance(selected_item['value'], str): + item_data = json.loads(selected_item['value']) + else: + item_data = selected_item['value'] + + logger.debug(f"Item data: {item_data}") + + item_id = item_data['id'] + return export_item_as_markdown(item_id) + except Exception as e: + error_message = f"Error processing selected item: {str(e)}" + logger.error(error_message) + return None, error_message + + search_results.select( + fn=handle_item_selection, + inputs=[search_results], + outputs=[export_output, error_output], + show_progress="full" + ) + + diff --git a/App_Function_Libraries/Gradio_UI/Gradio_Shared.py b/App_Function_Libraries/Gradio_UI/Gradio_Shared.py new file mode 100644 index 0000000000000000000000000000000000000000..83925ec9d41f0d68b90729cbdfd9aa7b83b7fb10 --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Gradio_Shared.py @@ -0,0 +1,285 @@ +# Gradio_Shared.py +# Gradio UI functions that are shared across multiple tabs +# +# Imports +import logging +import sqlite3 +import traceback +from functools import wraps +from typing import List, Tuple +# +# External Imports +import gradio as gr +# +# Local Imports +from App_Function_Libraries.DB.DB_Manager import list_prompts, db, search_and_display, fetch_prompt_details +from App_Function_Libraries.DB.SQLite_DB import DatabaseError +from App_Function_Libraries.Utils.Utils import format_transcription +# +############################################################################################################## +# +# Functions: + +whisper_models = ["small", "medium", "small.en", "medium.en", "medium", "large", "large-v1", "large-v2", "large-v3", + "distil-large-v2", "distil-medium.en", "distil-small.en"] + +# Sample data +prompts_category_1 = [ + "What are the key points discussed in the video?", + "Summarize the main arguments made by the speaker.", + "Describe the conclusions of the study presented." +] + +prompts_category_2 = [ + "How does the proposed solution address the problem?", + "What are the implications of the findings?", + "Can you explain the theory behind the observed phenomenon?" +] + +all_prompts = prompts_category_1 + prompts_category_2 + + + +#FIXME - SQL Functions that need to be addressed/added to DB manager +def search_media(query, fields, keyword, page): + try: + results = search_and_display(query, fields, keyword, page) + return results + except Exception as e: + logger = logging.getLogger() + logger.error(f"Error searching media: {e}") + return str(e) + +def fetch_items_by_title_or_url(search_query: str, search_type: str): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + if search_type == 'Title': + cursor.execute("SELECT id, title, url FROM Media WHERE title LIKE ?", (f'%{search_query}%',)) + elif search_type == 'URL': + cursor.execute("SELECT id, title, url FROM Media WHERE url LIKE ?", (f'%{search_query}%',)) + results = cursor.fetchall() + return results + except sqlite3.Error as e: + raise DatabaseError(f"Error fetching items by {search_type}: {e}") + +def fetch_items_by_keyword(search_query: str): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + SELECT m.id, m.title, m.url + FROM Media m + JOIN MediaKeywords mk ON m.id = mk.media_id + JOIN Keywords k ON mk.keyword_id = k.id + WHERE k.keyword LIKE ? + """, (f'%{search_query}%',)) + results = cursor.fetchall() + return results + except sqlite3.Error as e: + raise DatabaseError(f"Error fetching items by keyword: {e}") + +# FIXME - Raw SQL not using DB_Manager... +def fetch_items_by_content(search_query: str): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute("SELECT id, title, url FROM Media WHERE content LIKE ?", (f'%{search_query}%',)) + results = cursor.fetchall() + return results + except sqlite3.Error as e: + raise DatabaseError(f"Error fetching items by content: {e}") + + + +# FIXME - RAW SQL not using DB_Manager... +def fetch_item_details_single(media_id: int): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + SELECT prompt, summary + FROM MediaModifications + WHERE media_id = ? + ORDER BY modification_date DESC + LIMIT 1 + """, (media_id,)) + prompt_summary_result = cursor.fetchone() + cursor.execute("SELECT content FROM Media WHERE id = ?", (media_id,)) + content_result = cursor.fetchone() + + prompt = prompt_summary_result[0] if prompt_summary_result else "" + summary = prompt_summary_result[1] if prompt_summary_result else "" + content = content_result[0] if content_result else "" + + return prompt, summary, content + except sqlite3.Error as e: + raise Exception(f"Error fetching item details: {e}") + + +# FIXME - RAW SQL not using DB_Manager... +def fetch_item_details(media_id: int): + try: + with db.get_connection() as conn: + cursor = conn.cursor() + cursor.execute(""" + SELECT prompt, summary + FROM MediaModifications + WHERE media_id = ? + ORDER BY modification_date DESC + LIMIT 1 + """, (media_id,)) + prompt_summary_result = cursor.fetchone() + cursor.execute("SELECT content FROM Media WHERE id = ?", (media_id,)) + content_result = cursor.fetchone() + + prompt = prompt_summary_result[0] if prompt_summary_result else "" + summary = prompt_summary_result[1] if prompt_summary_result else "" + content = content_result[0] if content_result else "" + + return content, prompt, summary + except sqlite3.Error as e: + logging.error(f"Error fetching item details: {e}") + return "", "", "" # Return empty strings if there's an error + +# Handle prompt selection +def handle_prompt_selection(prompt): + return f"You selected: {prompt}" + + +def update_user_prompt(preset_name): + details = fetch_prompt_details(preset_name) + if details: + # Return a dictionary with all details + return { + "title": details[0], + "author": details[1], + "details": details[2], + "system_prompt": details[3], + "user_prompt": details[4] if len(details) > 3 else "", + } + return {"title": "", "details": "", "system_prompt": "", "user_prompt": "", "author": ""} + +def browse_items(search_query, search_type): + if search_type == 'Keyword': + results = fetch_items_by_keyword(search_query) + elif search_type == 'Content': + results = fetch_items_by_content(search_query) + else: + results = fetch_items_by_title_or_url(search_query, search_type) + return results + + +def update_dropdown(search_query, search_type): + results = browse_items(search_query, search_type) + item_options = [f"{item[1]} ({item[2]})" for item in results] + new_item_mapping = {f"{item[1]} ({item[2]})": item[0] for item in results} + print(f"Debug - Update Dropdown - New Item Mapping: {new_item_mapping}") + return gr.update(choices=item_options), new_item_mapping + + + +def get_media_id(selected_item, item_mapping): + return item_mapping.get(selected_item) + + +def update_detailed_view(item, item_mapping): + # Function to update the detailed view based on selected item + if item: + item_id = item_mapping.get(item) + if item_id: + content, prompt, summary = fetch_item_details(item_id) + if content or prompt or summary: + details_html = "

Details:

" + if prompt: + formatted_prompt = format_transcription(prompt) + details_html += f"

Prompt:

{formatted_prompt}

" + if summary: + formatted_summary = format_transcription(summary) + details_html += f"

Summary:

{formatted_summary}

" + # Format the transcription content for better readability + formatted_content = format_transcription(content) + #content_html = f"

Transcription:

{content}
" + content_html = f"

Transcription:

{formatted_content}
" + return details_html, content_html + else: + return "No details available.", "No details available." + else: + return "No item selected", "No item selected" + else: + return "No item selected", "No item selected" + + +def format_content(content): + # Format content using markdown + formatted_content = f"```\n{content}\n```" + return formatted_content + + +def update_prompt_dropdown(): + prompt_names = list_prompts() + return gr.update(choices=prompt_names) + + +def display_prompt_details(selected_prompt): + if selected_prompt: + prompts = update_user_prompt(selected_prompt) + if prompts["title"]: # Check if we have any details + details_str = f"

Details:

{prompts['details']}

" + system_str = f"

System:

{prompts['system_prompt']}

" + user_str = f"

User:

{prompts['user_prompt']}

" if prompts['user_prompt'] else "" + return details_str + system_str + user_str + return "No details available." + +def search_media_database(query: str) -> List[Tuple[int, str, str]]: + return browse_items(query, 'Title') + + +def load_media_content(media_id: int) -> dict: + try: + print(f"Debug - Load Media Content - Media ID: {media_id}") + item_details = fetch_item_details(media_id) + print(f"Debug - Load Media Content - Item Details: \n\n{item_details}\n\n\n\n") + + if isinstance(item_details, tuple) and len(item_details) == 3: + content, prompt, summary = item_details + else: + print(f"Debug - Load Media Content - Unexpected item_details format: \n\n{item_details}\n\n\n\n") + content, prompt, summary = "", "", "" + + return { + "content": content or "No content available", + "prompt": prompt or "No prompt available", + "summary": summary or "No summary available" + } + except Exception as e: + print(f"Debug - Load Media Content - Error: {str(e)}") + return {"content": "", "prompt": "", "summary": ""} + + +def error_handler(func): + @wraps(func) + def wrapper(*args, **kwargs): + try: + return func(*args, **kwargs) + except Exception as e: + error_message = f"Error in {func.__name__}: {str(e)}" + logging.error(f"{error_message}\n{traceback.format_exc()}") + return {"error": error_message, "details": traceback.format_exc()} + return wrapper + + +def create_chunking_inputs(): + chunk_text_by_words_checkbox = gr.Checkbox(label="Chunk Text by Words", value=False, visible=True) + max_words_input = gr.Number(label="Max Words", value=300, precision=0, visible=True) + chunk_text_by_sentences_checkbox = gr.Checkbox(label="Chunk Text by Sentences", value=False, visible=True) + max_sentences_input = gr.Number(label="Max Sentences", value=10, precision=0, visible=True) + chunk_text_by_paragraphs_checkbox = gr.Checkbox(label="Chunk Text by Paragraphs", value=False, visible=True) + max_paragraphs_input = gr.Number(label="Max Paragraphs", value=5, precision=0, visible=True) + chunk_text_by_tokens_checkbox = gr.Checkbox(label="Chunk Text by Tokens", value=False, visible=True) + max_tokens_input = gr.Number(label="Max Tokens", value=1000, precision=0, visible=True) + gr_semantic_chunk_long_file = gr.Checkbox(label="Semantic Chunking by Sentence similarity", value=False, visible=True) + gr_semantic_chunk_long_file_size = gr.Number(label="Max Chunk Size", value=2000, visible=True) + gr_semantic_chunk_long_file_overlap = gr.Number(label="Max Chunk Overlap Size", value=100, visible=True) + return [chunk_text_by_words_checkbox, max_words_input, chunk_text_by_sentences_checkbox, max_sentences_input, + chunk_text_by_paragraphs_checkbox, max_paragraphs_input, chunk_text_by_tokens_checkbox, max_tokens_input] diff --git a/App_Function_Libraries/Gradio_UI/Import_Functionality.py b/App_Function_Libraries/Gradio_UI/Import_Functionality.py new file mode 100644 index 0000000000000000000000000000000000000000..c748d2c866fc44f781a2a2e1c3045d7f4deff064 --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Import_Functionality.py @@ -0,0 +1,388 @@ +# Import_Functionality.py +# Functionality to import content into the DB +# +# Imports +from time import sleep +import logging +import re +import shutil +import tempfile +import os +import traceback +import zipfile +# +# External Imports +import gradio as gr +# +# Local Imports +from App_Function_Libraries.DB.DB_Manager import insert_prompt_to_db, load_preset_prompts, import_obsidian_note_to_db, \ + add_media_to_database +from App_Function_Libraries.Prompt_Handling import import_prompt_from_file, import_prompts_from_zip# +from App_Function_Libraries.Summarization.Summarization_General_Lib import perform_summarization + +################################################################################################################### +# +# Functions: + +logger = logging.getLogger() + + +def import_data(file, title, author, keywords, custom_prompt, summary, auto_summarize, api_name, api_key): + logging.debug(f"Starting import_data with file: {file} / Title: {title} / Author: {author} / Keywords: {keywords}") + if file is None: + return "No file uploaded. Please upload a file." + + try: + logging.debug(f"File object type: {type(file)}") + logging.debug(f"File object attributes: {dir(file)}") + + if hasattr(file, 'name'): + file_name = file.name + else: + file_name = 'unknown_file' + + # Create a temporary file + with tempfile.NamedTemporaryFile(mode='w+', delete=False, suffix='.txt', encoding='utf-8') as temp_file: + if isinstance(file, str): + # If file is a string, it's likely file content + temp_file.write(file) + elif hasattr(file, 'read'): + # If file has a 'read' method, it's likely a file-like object + content = file.read() + if isinstance(content, bytes): + content = content.decode('utf-8') + temp_file.write(content) + else: + # If it's neither a string nor a file-like object, try converting it to a string + temp_file.write(str(file)) + + temp_file.seek(0) + file_content = temp_file.read() + + logging.debug(f"File name: {file_name}") + logging.debug(f"File content (first 100 chars): {file_content[:100]}") + + # Create info_dict + info_dict = { + 'title': title or 'Untitled', + 'uploader': author or 'Unknown', + } + + # FIXME - Add chunking support... I added chapter chunking specifically for this... + # Create segments (assuming one segment for the entire content) + segments = [{'Text': file_content}] + + # Process keywords + keyword_list = [kw.strip() for kw in keywords.split(',') if kw.strip()] if keywords else [] + + # Handle summarization + if auto_summarize and api_name and api_key: + summary = perform_summarization(api_name, file_content, custom_prompt, api_key) + elif not summary: + summary = "No summary provided" + + # Add to database + result = add_media_to_database( + url=file_name, # Using filename as URL + info_dict=info_dict, + segments=segments, + summary=summary, + keywords=keyword_list, + custom_prompt_input=custom_prompt, + whisper_model="Imported", # Indicating this was an imported file + media_type="document", + overwrite=False # Set this to True if you want to overwrite existing entries + ) + + # Clean up the temporary file + os.unlink(temp_file.name) + + return f"File '{file_name}' import attempt complete. Database result: {result}" + except Exception as e: + logging.exception(f"Error importing file: {str(e)}") + return f"Error importing file: {str(e)}" + + +def process_obsidian_zip(zip_file): + with tempfile.TemporaryDirectory() as temp_dir: + try: + with zipfile.ZipFile(zip_file, 'r') as zip_ref: + zip_ref.extractall(temp_dir) + + imported_files, total_files, errors = import_obsidian_vault(temp_dir) + + return imported_files, total_files, errors + except zipfile.BadZipFile: + error_msg = "The uploaded file is not a valid zip file." + logger.error(error_msg) + return 0, 0, [error_msg] + except Exception as e: + error_msg = f"Error processing zip file: {str(e)}\n{traceback.format_exc()}" + logger.error(error_msg) + return 0, 0, [error_msg] + finally: + shutil.rmtree(temp_dir, ignore_errors=True) + + + +def scan_obsidian_vault(vault_path): + markdown_files = [] + for root, dirs, files in os.walk(vault_path): + for file in files: + if file.endswith('.md'): + markdown_files.append(os.path.join(root, file)) + return markdown_files + + +def parse_obsidian_note(file_path): + with open(file_path, 'r', encoding='utf-8') as file: + content = file.read() + + frontmatter = {} + frontmatter_match = re.match(r'^---\s*\n(.*?)\n---\s*\n', content, re.DOTALL) + if frontmatter_match: + frontmatter_text = frontmatter_match.group(1) + import yaml + frontmatter = yaml.safe_load(frontmatter_text) + content = content[frontmatter_match.end():] + + tags = re.findall(r'#(\w+)', content) + links = re.findall(r'\[\[(.*?)\]\]', content) + + return { + 'title': os.path.basename(file_path).replace('.md', ''), + 'content': content, + 'frontmatter': frontmatter, + 'tags': tags, + 'links': links, + 'file_path': file_path # Add this line + } + +def create_import_single_prompt_tab(): + with gr.TabItem("Import a Prompt", visible=True): + gr.Markdown("# Import a prompt into the database") + + with gr.Row(): + with gr.Column(): + import_file = gr.File(label="Upload file for import", file_types=["txt", "md"]) + title_input = gr.Textbox(label="Title", placeholder="Enter the title of the content") + author_input = gr.Textbox(label="Author", placeholder="Enter the author's name") + system_input = gr.Textbox(label="System", placeholder="Enter the system message for the prompt", lines=3) + user_input = gr.Textbox(label="User", placeholder="Enter the user message for the prompt", lines=3) + keywords_input = gr.Textbox(label="Keywords", placeholder="Enter keywords separated by commas") + import_button = gr.Button("Import Prompt") + + with gr.Column(): + import_output = gr.Textbox(label="Import Status") + save_button = gr.Button("Save to Database") + save_output = gr.Textbox(label="Save Status") + + def handle_import(file): + result = import_prompt_from_file(file) + if isinstance(result, tuple) and len(result) == 5: + title, author, system, user, keywords = result + return gr.update(value="File successfully imported. You can now edit the content before saving."), \ + gr.update(value=title), gr.update(value=author), gr.update(value=system), \ + gr.update(value=user), gr.update(value=", ".join(keywords)) + else: + return gr.update(value=result), gr.update(), gr.update(), gr.update(), gr.update(), gr.update() + + import_button.click( + fn=handle_import, + inputs=[import_file], + outputs=[import_output, title_input, author_input, system_input, user_input, keywords_input] + ) + + def save_prompt_to_db(title, author, system, user, keywords): + keyword_list = [k.strip() for k in keywords.split(',') if k.strip()] + return insert_prompt_to_db(title, author, system, user, keyword_list) + + save_button.click( + fn=save_prompt_to_db, + inputs=[title_input, author_input, system_input, user_input, keywords_input], + outputs=save_output + ) + + def update_prompt_dropdown(): + return gr.update(choices=load_preset_prompts()) + + save_button.click( + fn=update_prompt_dropdown, + inputs=[], + outputs=[gr.Dropdown(label="Select Preset Prompt")] + ) + +def create_import_item_tab(): + with gr.TabItem("Import Markdown/Text Files", visible=True): + gr.Markdown("# Import a markdown file or text file into the database") + gr.Markdown("...and have it tagged + summarized") + with gr.Row(): + with gr.Column(): + import_file = gr.File(label="Upload file for import", file_types=["txt", "md"]) + title_input = gr.Textbox(label="Title", placeholder="Enter the title of the content") + author_input = gr.Textbox(label="Author", placeholder="Enter the author's name") + keywords_input = gr.Textbox(label="Keywords", placeholder="Enter keywords, comma-separated") + custom_prompt_input = gr.Textbox(label="Custom Prompt", + placeholder="Enter a custom prompt for summarization (optional)") + summary_input = gr.Textbox(label="Summary", + placeholder="Enter a summary or leave blank for auto-summarization", lines=3) + auto_summarize_checkbox = gr.Checkbox(label="Auto-summarize", value=False) + api_name_input = gr.Dropdown( + choices=[None, "Local-LLM", "OpenAI", "Anthropic", "Cohere", "Groq", "DeepSeek", "Mistral", "OpenRouter", + "Llama.cpp", "Kobold", "Ooba", "Tabbyapi", "VLLM","ollama", "HuggingFace", "Custom-OpenAI-API"], + label="API for Auto-summarization" + ) + api_key_input = gr.Textbox(label="API Key", type="password") + with gr.Column(): + import_button = gr.Button("Import Data") + import_output = gr.Textbox(label="Import Status") + + import_button.click( + fn=import_data, + inputs=[import_file, title_input, author_input, keywords_input, custom_prompt_input, + summary_input, auto_summarize_checkbox, api_name_input, api_key_input], + outputs=import_output + ) + + +def create_import_multiple_prompts_tab(): + with gr.TabItem("Import Multiple Prompts", visible=True): + gr.Markdown("# Import multiple prompts into the database") + gr.Markdown("Upload a zip file containing multiple prompt files (txt or md)") + + with gr.Row(): + with gr.Column(): + zip_file = gr.File(label="Upload zip file for import", file_types=["zip"]) + import_button = gr.Button("Import Prompts") + prompts_dropdown = gr.Dropdown(label="Select Prompt to Edit", choices=[]) + title_input = gr.Textbox(label="Title", placeholder="Enter the title of the content") + author_input = gr.Textbox(label="Author", placeholder="Enter the author's name") + system_input = gr.Textbox(label="System", placeholder="Enter the system message for the prompt", + lines=3) + user_input = gr.Textbox(label="User", placeholder="Enter the user message for the prompt", lines=3) + keywords_input = gr.Textbox(label="Keywords", placeholder="Enter keywords separated by commas") + + with gr.Column(): + import_output = gr.Textbox(label="Import Status") + save_button = gr.Button("Save to Database") + save_output = gr.Textbox(label="Save Status") + prompts_display = gr.Textbox(label="Identified Prompts") + + def handle_zip_import(zip_file): + result = import_prompts_from_zip(zip_file) + if isinstance(result, list): + prompt_titles = [prompt['title'] for prompt in result] + return gr.update( + value="Zip file successfully imported. Select a prompt to edit from the dropdown."), prompt_titles, gr.update( + value="\n".join(prompt_titles)), result + else: + return gr.update(value=result), [], gr.update(value=""), [] + + def handle_prompt_selection(selected_title, prompts): + selected_prompt = next((prompt for prompt in prompts if prompt['title'] == selected_title), None) + if selected_prompt: + return ( + selected_prompt['title'], + selected_prompt.get('author', ''), + selected_prompt['system'], + selected_prompt.get('user', ''), + ", ".join(selected_prompt.get('keywords', [])) + ) + else: + return "", "", "", "", "" + + zip_import_state = gr.State([]) + + import_button.click( + fn=handle_zip_import, + inputs=[zip_file], + outputs=[import_output, prompts_dropdown, prompts_display, zip_import_state] + ) + + prompts_dropdown.change( + fn=handle_prompt_selection, + inputs=[prompts_dropdown, zip_import_state], + outputs=[title_input, author_input, system_input, user_input, keywords_input] + ) + + def save_prompt_to_db(title, author, system, user, keywords): + keyword_list = [k.strip() for k in keywords.split(',') if k.strip()] + return insert_prompt_to_db(title, author, system, user, keyword_list) + + save_button.click( + fn=save_prompt_to_db, + inputs=[title_input, author_input, system_input, user_input, keywords_input], + outputs=save_output + ) + + def update_prompt_dropdown(): + return gr.update(choices=load_preset_prompts()) + + save_button.click( + fn=update_prompt_dropdown, + inputs=[], + outputs=[gr.Dropdown(label="Select Preset Prompt")] + ) + + +def create_import_obsidian_vault_tab(): + with gr.TabItem("Import Obsidian Vault", visible=True): + gr.Markdown("## Import Obsidian Vault") + with gr.Row(): + with gr.Column(): + vault_path_input = gr.Textbox(label="Obsidian Vault Path (Local)") + vault_zip_input = gr.File(label="Upload Obsidian Vault (Zip)") + with gr.Column(): + import_vault_button = gr.Button("Import Obsidian Vault") + import_status = gr.Textbox(label="Import Status", interactive=False) + + + def import_vault(vault_path, vault_zip): + if vault_zip: + imported, total, errors = process_obsidian_zip(vault_zip.name) + elif vault_path: + imported, total, errors = import_obsidian_vault(vault_path) + else: + return "Please provide either a local vault path or upload a zip file." + + status = f"Imported {imported} out of {total} files.\n" + if errors: + status += f"Encountered {len(errors)} errors:\n" + "\n".join(errors) + return status + + + import_vault_button.click( + fn=import_vault, + inputs=[vault_path_input, vault_zip_input], + outputs=[import_status], + ) + + +def import_obsidian_vault(vault_path, progress=gr.Progress()): + try: + markdown_files = scan_obsidian_vault(vault_path) + total_files = len(markdown_files) + imported_files = 0 + errors = [] + + for i, file_path in enumerate(markdown_files): + try: + note_data = parse_obsidian_note(file_path) + success, error_msg = import_obsidian_note_to_db(note_data) + if success: + imported_files += 1 + else: + errors.append(error_msg) + except Exception as e: + error_msg = f"Error processing {file_path}: {str(e)}" + logger.error(error_msg) + errors.append(error_msg) + + progress((i + 1) / total_files, f"Imported {imported_files} of {total_files} files") + sleep(0.1) # Small delay to prevent UI freezing + + return imported_files, total_files, errors + except Exception as e: + error_msg = f"Error scanning vault: {str(e)}\n{traceback.format_exc()}" + logger.error(error_msg) + return 0, 0, [error_msg] \ No newline at end of file diff --git a/App_Function_Libraries/Gradio_UI/Introduction_tab.py b/App_Function_Libraries/Gradio_UI/Introduction_tab.py new file mode 100644 index 0000000000000000000000000000000000000000..9942a89f81b6f3fce4849cd5a57790a15bd90d5e --- /dev/null +++ b/App_Function_Libraries/Gradio_UI/Introduction_tab.py @@ -0,0 +1,167 @@ +# Introduction_tab.py +# Gradio UI functions for the Introduction tab +# +# Imports +# +# External Imports +import gradio as gr +# +# Local Imports +from App_Function_Libraries.DB.DB_Manager import get_db_config +# +#################################################################################################### +# +# Functions: + + + +def create_introduction_tab(): + with gr.TabItem("Introduction", visible=True): + db_config = get_db_config() + db_type = db_config['type'] + gr.Markdown(f"# tldw: Your LLM-powered Research Multi-tool (Using {db_type.capitalize()} Database)") + with gr.Row(): + with gr.Column(): + gr.Markdown("""### What can it do? + - Transcribe and summarize videos from URLs/Local files + - Transcribe and Summarize Audio files/Podcasts (URL/local file) + - Summarize articles from URLs/Local notes + - Ingest and summarize books(epub/PDF) + - Ingest and summarize research papers (PDFs - WIP) + - Search and display ingested content + summaries + - Create and manage custom prompts + - Chat with an LLM of your choice to generate content using the selected item + Prompts + - Keyword support for content search and display + - Export keywords/items to markdown/CSV(csv is wip) + - Import existing notes from Obsidian to the database (Markdown/txt files or a zip containing a collection of files) + - View and manage chat history + - Writing Tools: Grammar & Style check, Tone Analyzer & Editor, more planned... + - RAG (Retrieval-Augmented Generation) support for content generation(think about asking questions about your entire library of items) + - More features planned... + - All powered by your choice of LLM. + - Currently supports: Local-LLM(llamafile-server), OpenAI, Anthropic, Cohere, Groq, DeepSeek, OpenRouter, Llama.cpp, Kobold, Ooba, Tabbyapi, VLLM and more to come... + - All data is stored locally in a SQLite database for easy access and management. + - No trackers (Gradio has some analytics but it's disabled here...) + - No ads, no tracking, no BS. Just you and your content. + - Open-source and free to use. Contributions welcome! + - If you have any thoughts or feedback, please let me know on github or via email. + """) + gr.Markdown( + """Follow this project at [tl/dw: Too Long, Didn't Watch - Your Personal Research Multi-Tool - GitHub](https://github.com/rmusser01/tldw)""") + with gr.Column(): + gr.Markdown("""### How to use: + ##### Quick Start: Just click on the appropriate tab for what you're trying to do and fill in the required fields. Click "Process