File size: 37,472 Bytes
f4368a7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
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"<span style='color: red'>โณ {status}</span>"
        elif status == "completed":
            status_formatted = f"<span style='color: green'>โœ… {status}</span>"
        else:
            status_formatted = f"<span style='color: orange'>โ“ {status}</span>"
        
        # 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)