import os os.environ["TOKENIZERS_PARALLELISM"] = "false" import datetime import functools import traceback from typing import List, Optional, Any, Dict import torch import transformers from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline from langchain_community.llms import HuggingFacePipeline # Other LangChain and community imports from langchain_community.document_loaders import OnlinePDFLoader from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain_community.vectorstores import FAISS from langchain.embeddings import HuggingFaceEmbeddings from langchain_community.retrievers import BM25Retriever from langchain.retrievers import EnsembleRetriever from langchain.prompts import ChatPromptTemplate from langchain.schema import StrOutputParser, Document from langchain_core.runnables import RunnableParallel, RunnableLambda from transformers.quantizers.auto import AutoQuantizationConfig import gradio as gr import requests from pydantic import PrivateAttr import pydantic from langchain.llms.base import LLM from typing import Any, Optional, List import typing import time print("Pydantic Version: ") print(pydantic.__version__) # Add Mistral imports with fallback handling try: from mistralai import Mistral MISTRAL_AVAILABLE = True debug_print = lambda msg: print(f"[{datetime.datetime.now().isoformat()}] {msg}") debug_print("Loaded latest Mistral client library") except ImportError: MISTRAL_AVAILABLE = False debug_print = lambda msg: print(f"[{datetime.datetime.now().isoformat()}] {msg}") debug_print("Mistral client library not found. Install with: pip install mistralai") def debug_print(message: str): print(f"[{datetime.datetime.now().isoformat()}] {message}", flush=True) def word_count(text: str) -> int: return len(text.split()) # Initialize a tokenizer for token counting (using gpt2 as a generic fallback) def initialize_tokenizer(): try: return AutoTokenizer.from_pretrained("gpt2") except Exception as e: debug_print("Failed to initialize tokenizer: " + str(e)) return None global_tokenizer = initialize_tokenizer() def count_tokens(text: str) -> int: if global_tokenizer: try: return len(global_tokenizer.encode(text)) except Exception as e: return len(text.split()) return len(text.split()) # Add these imports at the top of your file import uuid import threading import queue from typing import Dict, Any, Tuple, Optional import time # Global storage for jobs and results jobs = {} # Stores job status and results results_queue = queue.Queue() # Thread-safe queue for completed jobs processing_lock = threading.Lock() # Prevent simultaneous processing of the same job # Add a global variable to store the last job ID last_job_id = None # Add these missing async processing functions def process_in_background(job_id, function, args): """Process a function in the background and store results""" try: debug_print(f"Processing job {job_id} in background") result = function(*args) results_queue.put((job_id, result)) debug_print(f"Job {job_id} completed and added to results queue") except Exception as e: debug_print(f"Error in background job {job_id}: {str(e)}") error_result = (f"Error processing job: {str(e)}", "", "", "") results_queue.put((job_id, error_result)) def load_pdfs_async(file_links, model_choice, prompt_template, bm25_weight, temperature, top_p): """Asynchronous version of load_pdfs_updated to prevent timeouts""" global last_job_id if not file_links: return "Please enter non-empty URLs", "", "Model used: N/A", "", "", get_job_list() job_id = str(uuid.uuid4()) debug_print(f"Starting async job {job_id} for file loading") # Start background thread threading.Thread( target=process_in_background, args=(job_id, load_pdfs_updated, [file_links, model_choice, prompt_template, bm25_weight, temperature, top_p]) ).start() job_query = f"Loading files: {file_links.split()[0]}..." if file_links else "No files" jobs[job_id] = { "status": "processing", "type": "load_files", "start_time": time.time(), "query": job_query } last_job_id = job_id return ( f"Files submitted and processing in the background (Job ID: {job_id}).\n\n" f"Use 'Check Job Status' tab with this ID to get results.", f"Job ID: {job_id}", f"Model requested: {model_choice}", job_id, # Return job_id to update the job_id_input component job_query, # Return job_query to update the job_query_display component get_job_list() # Return updated job list ) def submit_query_async(query, use_llama, use_mistral, temperature, top_p): """Asynchronous version of submit_query_updated to prevent timeouts""" global last_job_id if not query: return ("Please enter a non-empty query", "Input/Output tokens: 0/0", "Please enter a non-empty query", "Input/Output tokens: 0/0", "", "", get_job_list()) if not (use_llama or use_mistral): return ("Please select at least one model", "Input/Output tokens: 0/0", "Please select at least one model", "Input/Output tokens: 0/0", "", "", get_job_list()) responses = {"llama": None, "mistral": None} job_ids = [] if use_llama: llama_job_id = str(uuid.uuid4()) debug_print(f"Starting async job {llama_job_id} for Llama query: {query}") # Start background thread for Llama threading.Thread( target=process_in_background, args=(llama_job_id, submit_query_updated, [query, "πŸ‡ΊπŸ‡Έ Remote Meta-Llama-3", temperature, top_p]) ).start() jobs[llama_job_id] = { "status": "processing", "type": "query", "start_time": time.time(), "query": query, "model": "Llama" } job_ids.append(llama_job_id) responses["llama"] = f"Processing (Job ID: {llama_job_id})" if use_mistral: mistral_job_id = str(uuid.uuid4()) debug_print(f"Starting async job {mistral_job_id} for Mistral query: {query}") # Start background thread for Mistral threading.Thread( target=process_in_background, args=(mistral_job_id, submit_query_updated, [query, "πŸ‡ͺπŸ‡Ί Mistral-API", temperature, top_p]) ).start() jobs[mistral_job_id] = { "status": "processing", "type": "query", "start_time": time.time(), "query": query, "model": "Mistral" } job_ids.append(mistral_job_id) responses["mistral"] = f"Processing (Job ID: {mistral_job_id})" # Store the last job ID (use the first one for now) last_job_id = job_ids[0] if job_ids else None return ( responses.get("llama", "Not selected"), "Input tokens: " + str(count_tokens(query)) if use_llama else "Not selected", responses.get("mistral", "Not selected"), "Input tokens: " + str(count_tokens(query)) if use_mistral else "Not selected", last_job_id, query, get_job_list() ) def update_ui_with_last_job_id(): # This function doesn't need to do anything anymore # We'll update the UI directly in the functions that call this pass # Function to display all jobs as a clickable list def get_job_list(): job_list_md = "### Submitted Jobs\n\n" if not jobs: return "No jobs found. Submit a query or load files to create jobs." # Sort jobs by start time (newest first) sorted_jobs = sorted( [(job_id, job_info) for job_id, job_info in jobs.items()], key=lambda x: x[1].get("start_time", 0), reverse=True ) for job_id, job_info in sorted_jobs: status = job_info.get("status", "unknown") job_type = job_info.get("type", "unknown") query = job_info.get("query", "") model = job_info.get("model", "") # Get the model name start_time = job_info.get("start_time", 0) time_str = datetime.datetime.fromtimestamp(start_time).strftime("%Y-%m-%d %H:%M:%S") # Create a shortened query preview query_preview = query[:30] + "..." if query and len(query) > 30 else query or "N/A" # Add color and icons based on status if status == "processing": status_formatted = f"⏳ {status}" elif status == "completed": status_formatted = f"βœ… {status}" else: status_formatted = f"❓ {status}" # Add model icon based on model name model_icon = "πŸ‡ΊπŸ‡Έ" if model == "Llama" else "πŸ‡ͺπŸ‡Ί" if model == "Mistral" else "" model_prefix = f"{model_icon} {model} " if model else "" # Create clickable links using Markdown if job_type == "query": job_list_md += f"- [{job_id}](javascript:void) - {time_str} - {status_formatted} - {model_prefix}Query: {query_preview}\n" else: job_list_md += f"- [{job_id}](javascript:void) - {time_str} - {status_formatted} - File Load Job\n" return job_list_md # Function to handle job list clicks def job_selected(job_id): if job_id in jobs: return job_id, jobs[job_id].get("query", "No query for this job") return job_id, "Job not found" # Function to refresh the job list def refresh_job_list(): return get_job_list() # Function to sync model dropdown boxes def sync_model_dropdown(value): return value # Function to check job status def check_job_status(job_id): if not job_id: return "Please enter a job ID", "", "", "", "" # Process any completed jobs in the queue try: while not results_queue.empty(): completed_id, result = results_queue.get_nowait() if completed_id in jobs: jobs[completed_id]["status"] = "completed" jobs[completed_id]["result"] = result jobs[completed_id]["end_time"] = time.time() debug_print(f"Job {completed_id} completed and stored in jobs dictionary") except queue.Empty: pass # Check if the requested job exists if job_id not in jobs: return "Job not found. Please check the ID and try again.", "", "", "", "" job = jobs[job_id] job_query = job.get("query", "No query available for this job") # If job is still processing if job["status"] == "processing": elapsed_time = time.time() - job["start_time"] job_type = job.get("type", "unknown") if job_type == "load_files": return ( f"Files are still being processed (elapsed: {elapsed_time:.1f}s).\n\n" f"Try checking again in a few seconds.", f"Job ID: {job_id}", f"Status: Processing", "", job_query ) else: # query job return ( f"Query is still being processed (elapsed: {elapsed_time:.1f}s).\n\n" f"Try checking again in a few seconds.", f"Job ID: {job_id}", f"Input tokens: {count_tokens(job.get('query', ''))}", "Output tokens: pending", job_query ) # If job is completed if job["status"] == "completed": result = job["result"] processing_time = job["end_time"] - job["start_time"] if job.get("type") == "load_files": return ( f"{result[0]}\n\nProcessing time: {processing_time:.1f}s", result[1], result[2], "", job_query ) else: # query job return ( f"{result[0]}\n\nProcessing time: {processing_time:.1f}s", result[1], result[2], result[3], job_query ) # Fallback for unknown status return f"Job status: {job['status']}", "", "", "", job_query # Function to clean up old jobs def cleanup_old_jobs(): current_time = time.time() to_delete = [] for job_id, job in jobs.items(): # Keep completed jobs for 24 hours, processing jobs for 48 hours if job["status"] == "completed" and (current_time - job.get("end_time", 0)) > 86400: to_delete.append(job_id) elif job["status"] == "processing" and (current_time - job.get("start_time", 0)) > 172800: to_delete.append(job_id) for job_id in to_delete: del jobs[job_id] debug_print(f"Cleaned up {len(to_delete)} old jobs. {len(jobs)} jobs remaining.") return f"Cleaned up {len(to_delete)} old jobs", "", "" # Improve the truncate_prompt function to be more aggressive with limiting context def truncate_prompt(prompt: str, max_tokens: int = 4096) -> str: """Truncate prompt to fit within token limit, preserving the most recent/relevant parts.""" if not prompt: return "" if global_tokenizer: try: tokens = global_tokenizer.encode(prompt) if len(tokens) > max_tokens: # For prompts, we often want to keep the beginning instructions and the end context # So we'll keep the first 20% and the last 80% of the max tokens beginning_tokens = int(max_tokens * 0.2) ending_tokens = max_tokens - beginning_tokens new_tokens = tokens[:beginning_tokens] + tokens[-(ending_tokens):] return global_tokenizer.decode(new_tokens) except Exception as e: debug_print(f"Truncation error: {str(e)}") # Fallback to word-based truncation words = prompt.split() if len(words) > max_tokens: beginning_words = int(max_tokens * 0.2) ending_words = max_tokens - beginning_words return " ".join(words[:beginning_words] + words[-(ending_words):]) return prompt default_prompt = """\ {conversation_history} Use the following context to provide a detailed technical answer to the user's question. Do not include an introduction like "Based on the provided documents, ...". Just answer the question. If you don't know the answer, please respond with "I don't know". Context: {context} User's question: {question} """ def load_txt_from_url(url: str) -> Document: response = requests.get(url) if response.status_code == 200: text = response.text.strip() if not text: raise ValueError(f"TXT file at {url} is empty.") return Document(page_content=text, metadata={"source": url}) else: raise Exception(f"Failed to load {url} with status {response.status_code}") class RemoteLLM(LLM): temperature: float = 0.5 top_p: float = 0.95 def __init__(self, temperature: float = 0.5, top_p: float = 0.95): super().__init__() self.temperature = temperature self.top_p = top_p @property def _llm_type(self) -> str: return "remote_llm" def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str: try: response = requests.post( "http://localhost:11434/api/generate", json={ "model": "llama2", "prompt": prompt, "temperature": self.temperature, "top_p": self.top_p }, stream=False ) if response.status_code == 200: return response.json()["response"] else: return f"Error: {response.status_code}" except Exception as e: return f"Error: {str(e)}" @property def _identifying_params(self) -> dict: return { "temperature": self.temperature, "top_p": self.top_p } class MistralLLM(LLM): temperature: float = 0.7 top_p: float = 0.95 _client: Any = PrivateAttr(default=None) def __init__(self, api_key: str, temperature: float = 0.7, top_p: float = 0.95, **kwargs: Any): try: super().__init__(**kwargs) object.__setattr__(self, '_client', Mistral(api_key=api_key)) self.temperature = temperature self.top_p = top_p except Exception as e: debug_print(f"Init Mistral failed with error: {e}") @property def _llm_type(self) -> str: return "mistral_llm" def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str: try: debug_print("Calling Mistral API...") response = self._client.chat.complete( model="mistral-small-latest", messages=[{"role": "user", "content": prompt}], temperature=self.temperature, top_p=self.top_p ) return response.choices[0].message.content except Exception as e: debug_print(f"Mistral API error: {str(e)}") return f"Error generating response: {str(e)}" @property def _identifying_params(self) -> dict: return {"model": "mistral-small-latest"} class LocalLLM(LLM): @property def _llm_type(self) -> str: return "local_llm" def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str: truncated_prompt = truncate_prompt(prompt) return f"Local LLM Fallback Response for: {truncated_prompt[:100]}..." @property def _identifying_params(self) -> dict: return {} class ErrorLLM(LLM): @property def _llm_type(self) -> str: return "error_llm" def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str: return "Error: LLM pipeline could not be created. Please check your configuration and try again." @property def _identifying_params(self) -> dict: return {} class SimpleLLMChain: def __init__(self, llm_choice: str = "Meta-Llama-3", temperature: float = 0.5, top_p: float = 0.95) -> None: self.llm_choice = llm_choice self.temperature = temperature self.top_p = top_p self.llm = self.create_llm_pipeline() self.conversation_history = [] # Keep track of conversation def create_llm_pipeline(self): from langchain.llms.base import LLM # Import LLM here so it's always defined normalized = self.llm_choice.lower() try: if "remote" in normalized: debug_print("Creating remote Meta-Llama-3 pipeline via Hugging Face Inference API...") from huggingface_hub import InferenceClient repo_id = "meta-llama/Meta-Llama-3-8B-Instruct" hf_api_token = os.environ.get("HF_API_TOKEN") if not hf_api_token: raise ValueError("Please set the HF_API_TOKEN environment variable to use remote inference.") client = InferenceClient(token=hf_api_token, timeout=120) # We no longer use wait_for_model because it's unsupported def remote_generate(prompt: str) -> str: max_retries = 3 backoff = 2 # start with 2 seconds for attempt in range(max_retries): try: debug_print(f"Remote generation attempt {attempt+1}") response = client.text_generation( prompt, model=repo_id, temperature=self.temperature, top_p=self.top_p, max_new_tokens=512 # Reduced token count for speed ) return response except Exception as e: debug_print(f"Attempt {attempt+1} failed with error: {e}") if attempt == max_retries - 1: raise time.sleep(backoff) backoff *= 2 # exponential backoff return "Failed to generate response after multiple attempts." class RemoteLLM(LLM): @property def _llm_type(self) -> str: return "remote_llm" def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str: return remote_generate(prompt) @property def _identifying_params(self) -> dict: return {"model": repo_id} debug_print("Remote Meta-Llama-3 pipeline created successfully.") return RemoteLLM() elif "mistral" in normalized: api_key = os.getenv("MISTRAL_API_KEY") return MistralLLM(api_key=api_key, temperature=self.temperature, top_p=self.top_p) else: return LocalLLM() except Exception as e: debug_print(f"Error creating LLM pipeline: {str(e)}") return ErrorLLM() def update_llm_pipeline(self, new_model_choice: str, temperature: float, top_p: float): self.llm_choice = new_model_choice self.temperature = temperature self.top_p = top_p self.llm = self.create_llm_pipeline() def submit_query(self, query: str) -> tuple: try: response = self.llm(query) # Store in conversation history self.conversation_history.append({"query": query, "response": response}) input_tokens = count_tokens(query) output_tokens = count_tokens(response) return (response, f"Input tokens: {input_tokens}", f"Output tokens: {output_tokens}") except Exception as e: return (f"Error processing query: {str(e)}", "Input tokens: 0", "Output tokens: 0") # Update submit_query_updated to work with the simplified chain def submit_query_updated(query: str, model_choice: str = None, temperature: float = 0.5, top_p: float = 0.95): """Process a query with the specified model and parameters.""" debug_print(f"Processing query: {query}") if not query: debug_print("Empty query received") return "Please enter a non-empty query", "", "Input tokens: 0", "Output tokens: 0" try: global llm_chain if llm_chain is None: llm_chain = SimpleLLMChain( llm_choice=model_choice, temperature=temperature, top_p=top_p ) elif llm_chain.llm_choice != model_choice: llm_chain.update_llm_pipeline(model_choice, temperature, top_p) response, input_tokens, output_tokens = llm_chain.submit_query(query) return response, "", input_tokens, output_tokens except Exception as e: debug_print(f"Error in submit_query_updated: {str(e)}") return f"Error: {str(e)}", "", "Input tokens: 0", "Output tokens: 0" # Update the reset_app_updated function def reset_app_updated(): global llm_chain llm_chain = None return "Application reset successfully" # ---------------------------- # Gradio Interface Functions # ---------------------------- global rag_chain rag_chain = SimpleLLMChain() def load_pdfs_updated(file_links, model_choice, prompt_template, bm25_weight, temperature, top_p): debug_print("Inside load_pdfs function.") if not file_links: debug_print("Please enter non-empty URLs") return "Please enter non-empty URLs", "Word count: N/A", "Model used: N/A", "Context: N/A" try: links = [link.strip() for link in file_links.split("\n") if link.strip()] global rag_chain if rag_chain.raw_data: rag_chain.update_llm_pipeline(model_choice, temperature, top_p, prompt_template, bm25_weight) context_display = rag_chain.get_current_context() response_msg = f"Files already loaded. Chain updated with model: {model_choice}" return ( response_msg, f"Word count: {word_count(rag_chain.context)}", f"Model used: {rag_chain.llm_choice}", f"Context:\n{context_display}" ) else: rag_chain = SimpleLLMChain( llm_choice=model_choice, temperature=temperature, top_p=top_p ) rag_chain.add_pdfs_to_vectore_store(links) context_display = rag_chain.get_current_context() response_msg = f"Files loaded successfully. Using model: {model_choice}" return ( response_msg, f"Word count: {word_count(rag_chain.context)}", f"Model used: {rag_chain.llm_choice}", f"Context:\n{context_display}" ) except Exception as e: error_msg = traceback.format_exc() debug_print("Could not load files. Error: " + error_msg) return ( "Error loading files: " + str(e), f"Word count: {word_count('')}", f"Model used: {rag_chain.llm_choice}", "Context: N/A" ) def update_model(new_model: str): global rag_chain if rag_chain and rag_chain.raw_data: rag_chain.update_llm_pipeline(new_model, rag_chain.temperature, rag_chain.top_p, rag_chain.prompt_template, rag_chain.bm25_weight) debug_print(f"Model updated to {rag_chain.llm_choice}") return f"Model updated to: {rag_chain.llm_choice}" else: return "No files loaded; please load files first." def reset_app_updated(): global rag_chain rag_chain = SimpleLLMChain() debug_print("App reset successfully.") return ( "App reset successfully. You can now load new files", "", "Model used: Not selected" ) # ---------------------------- # Gradio Interface Setup # ---------------------------- custom_css = """ textarea { overflow-y: scroll !important; max-height: 200px; } """ # Function to add dots and reset def add_dots_and_reset(): if not hasattr(add_dots_and_reset, "dots"): add_dots_and_reset.dots = "" # Initialize the attribute # Add a dot add_dots_and_reset.dots += "." # Reset after 5 dots if len(add_dots_and_reset.dots) > 5: add_dots_and_reset.dots = "" print(f"Current dots: {add_dots_and_reset.dots}") # Debugging print return add_dots_and_reset.dots # Define a dummy function to simulate data retrieval def run_query(max_value): # Simulate a data retrieval or processing function return [[i, i**2] for i in range(1, max_value + 1)] # Function to call both refresh_job_list and check_job_status using the last job ID def periodic_update(is_checked): interval = 2 if is_checked else None debug_print(f"Auto-refresh checkbox is {'checked' if is_checked else 'unchecked'}, every={interval}") if is_checked: global last_job_id job_list_md = refresh_job_list() job_status = check_job_status(last_job_id) if last_job_id else ("No job ID available", "", "", "", "") query_results = run_query(10) # Use a fixed value or another logic if needed return job_list_md, job_status[0], query_results, "" # Return empty string instead of context else: # Return empty values to stop updates return "", "", [], "" # Define a function to determine the interval based on the checkbox state def get_interval(is_checked): return 2 if is_checked else None # Update the Gradio interface to include job status checking with gr.Blocks(css=custom_css, js=""" document.addEventListener('DOMContentLoaded', function() { // Add event listener for job list clicks const jobListInterval = setInterval(() => { const jobLinks = document.querySelectorAll('.job-list-container a'); if (jobLinks.length > 0) { jobLinks.forEach(link => { link.addEventListener('click', function(e) { e.preventDefault(); const jobId = this.textContent.split(' ')[0]; // Find the job ID input textbox and set its value const jobIdInput = document.querySelector('.job-id-input input'); if (jobIdInput) { jobIdInput.value = jobId; // Trigger the input event to update Gradio's state jobIdInput.dispatchEvent(new Event('input', { bubbles: true })); } }); }); clearInterval(jobListInterval); } }, 500); }); """) as app: gr.Markdown('''# PsyLLM Interface **Model Selection & Parameters:** Choose from the following options: - πŸ‡ΊπŸ‡Έ Remote Meta-Llama-3 - has context windows of 8000 tokens - πŸ‡ͺπŸ‡Ί Mistral-API - has context windows of 32000 tokens **πŸ”₯ Randomness (Temperature):** Adjusts output predictability. - Example: 0.2 makes the output very deterministic (less creative), while 0.8 introduces more variety and spontaneity. **🎯 Word Variety (Top‑p):** Limits word choices to a set probability percentage. - Example: 0.5 restricts output to the most likely 50% of token choices for a focused answer; 0.95 allows almost all possibilities for more diverse responses. **⚠️ IMPORTANT: This app uses asynchronous processing to avoid timeout issues** - When you submit a query, you'll receive a Job ID - Use the "Check Job Status" tab to monitor and retrieve your results ''') with gr.Tabs() as tabs: with gr.TabItem("Submit Query"): with gr.Row(): with gr.Column(scale=1): llama_checkbox = gr.Checkbox( value=True, label="πŸ‡ΊπŸ‡Έ Remote Meta-Llama-3", info="Context window: 8000 tokens" ) mistral_checkbox = gr.Checkbox( value=False, label="πŸ‡ͺπŸ‡Ί Mistral-API", info="Context window: 32000 tokens" ) with gr.Column(scale=2): temperature_slider = gr.Slider( minimum=0.1, maximum=1.0, value=0.5, step=0.1, label="Randomness (Temperature)" ) top_p_slider = gr.Slider( minimum=0.1, maximum=0.99, value=0.95, step=0.05, label="Word Variety (Top-p)" ) with gr.Row(): query_input = gr.Textbox( label="Enter your query here", placeholder="Type your query", lines=4 ) submit_button = gr.Button("Submit Query to Selected Models") with gr.Row(): with gr.Column(scale=1): gr.Markdown("### Llama Results") llama_response = gr.Textbox( label="Llama Response", placeholder="Response will appear here", lines=6 ) llama_tokens = gr.Markdown("Input/Output tokens: 0/0") with gr.Column(scale=1): gr.Markdown("### Mistral Results") mistral_response = gr.Textbox( label="Mistral Response", placeholder="Response will appear here", lines=6 ) mistral_tokens = gr.Markdown("Input/Output tokens: 0/0") with gr.TabItem("Check Job Status"): with gr.Row(): with gr.Column(scale=1): job_list = gr.Markdown( value="No jobs yet", label="Job List (Click to select)" ) # Add the Refresh Job List button refresh_button = gr.Button("Refresh Job List") # Use a Checkbox to control the periodic updates auto_refresh_checkbox = gr.Checkbox( label="Enable Auto Refresh", value=False # Default to unchecked ) # Use a DataFrame to display results df = gr.DataFrame( value=run_query(10), # Initial value headers=["Number", "Square"], label="Query Results", visible=False # Set the DataFrame to be invisible ) with gr.Column(scale=2): job_id_input = gr.Textbox( label="Job ID", placeholder="Job ID will appear here when selected from the list", lines=1 ) job_query_display = gr.Textbox( label="Job Query", placeholder="The query associated with this job will appear here", lines=2, interactive=False ) check_button = gr.Button("Check Status") cleanup_button = gr.Button("Cleanup Old Jobs") with gr.Row(): status_response = gr.Textbox( label="Job Result", placeholder="Job result will appear here", lines=6 ) status_context = gr.Textbox( label="Context Information", placeholder="Context information will appear here", lines=6 ) with gr.Row(): status_tokens1 = gr.Markdown("") status_tokens2 = gr.Markdown("") with gr.TabItem("App Management"): with gr.Row(): reset_button = gr.Button("Reset App") with gr.Row(): reset_response = gr.Textbox( label="Reset Response", placeholder="Reset confirmation will appear here", lines=2 ) reset_context = gr.Textbox( label="", placeholder="", lines=2, visible=False ) with gr.Row(): reset_model = gr.Markdown("") # Connect the buttons to their respective functions submit_button.click( submit_query_async, inputs=[ query_input, llama_checkbox, mistral_checkbox, temperature_slider, top_p_slider ], outputs=[ llama_response, llama_tokens, mistral_response, mistral_tokens, job_id_input, job_query_display, job_list ] ) check_button.click( check_job_status, inputs=[job_id_input], outputs=[status_response, status_context, status_tokens1, status_tokens2, job_query_display] ) refresh_button.click( refresh_job_list, inputs=[], outputs=[job_list] ) job_id_input.change( job_selected, inputs=[job_id_input], outputs=[job_id_input, job_query_display] ) cleanup_button.click( cleanup_old_jobs, inputs=[], outputs=[status_response, status_context, status_tokens1] ) reset_button.click( reset_app_updated, inputs=[], outputs=[reset_response, reset_context, reset_model] ) app.load( fn=refresh_job_list, inputs=None, outputs=job_list ) auto_refresh_checkbox.change( fn=periodic_update, inputs=[auto_refresh_checkbox], outputs=[job_list, status_response, df, status_context], every=2 ) # Add this with your other global variables global llm_chain llm_chain = None if __name__ == "__main__": debug_print("Launching Gradio interface.") app.queue().launch(share=False)