-
Notifications
You must be signed in to change notification settings - Fork 79
/
controller.py
2722 lines (2302 loc) · 124 KB
/
controller.py
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
994
995
996
997
998
999
1000
import functools
import os
import tkinter as tk
import threading
from collections import defaultdict, ChainMap
import time
from functools import reduce
from pprint import pprint
from tkinter import filedialog, ttk
from tkinter import messagebox
from tkinter.font import Font
import re
import datetime
import PIL
import pyperclip
import bisect
import traceback
from view.colors import history_color, not_visited_color, visited_color, ooc_color, text_color, uncanonical_color, \
immutable_color
from view.display import Display
from components.dialogs import *
from model import TreeModel
from util.util import clip_num, metadata, diff, split_indices, diff_linesToWords
from util.util_tree import ancestry_in_range, depth, height, flatten_tree, stochastic_transition, node_ancestry, subtree_list, \
node_index, nearest_common_ancestor, filtered_children
from util.gpt_util import logprobs_to_probs, parse_logit_bias
from util.textbox_util import distribute_textbox_changes
from util.keybindings import tkinter_keybindings
from view.icons import Icons
from difflib import SequenceMatcher
from diff_match_patch import diff_match_patch
import json
from view.colors import edit_color, bg_color
def gated_call(f, condition):
def _gated_call(*args, _f=f, _cond=condition, **kwargs):
if _cond():
# print('cond')
_f(*args, **kwargs)
# print("no cond")
return _gated_call
def no_junk_args(f):
return lambda event=None, *args, _f=f, **kwargs: _f(*args, **kwargs)
class Controller:
def __init__(self, root):
self.callbacks = self.build_callbacks()
self.root = root
self.state = TreeModel(self.root)
self.display = Display(self.root, self.callbacks, self.state, self)
self.icons = Icons()
self.register_model_callbacks()
self.setup_key_bindings()
self.build_menus()
self.ancestor_end_indices = None
self.nav_history = []
self.undo_history = []
#################################
# Hooks
#################################
def build_callbacks(self):
attrs = [getattr(self, f) for f in dir(self)]
callbacks = [f for f in attrs if callable(f) and hasattr(f, "meta") and "name" in f.meta]
return {
f.meta["name"]: {**f.meta, "callback": no_junk_args(f)}
for f in callbacks
}
def register_model_callbacks(self):
# When the tree is updated, refresh the navtree, nav selection, and textbox
self.state.register_callback(self.state.tree_updated, self.update_nav_tree)
self.state.register_callback(self.state.tree_updated, self.fix_selection)
self.state.register_callback(self.state.tree_updated, self.update_nav_tree_selected)
self.state.register_callback(self.state.tree_updated, self.update_chapter_nav_tree)
self.state.register_callback(self.state.tree_updated, self.update_chapter_nav_tree_selected)
# TODO save_edits causes tree_updated...
self.state.register_callback(self.state.tree_updated, self.save_edits)
self.state.register_callback(self.state.tree_updated, self.refresh_textbox)
self.state.register_callback(self.state.tree_updated, self.refresh_visualization)
self.state.register_callback(self.state.tree_updated, self.refresh_display)
self.state.register_callback(self.state.tree_updated, self.setup_custom_key_bindings)
self.state.register_callback(self.state.tree_updated, self.modules_tree_updated)
# TODO autosaving takes too long for a big tree
self.state.register_callback(self.state.tree_updated, lambda **kwargs: self.save_tree(popup=False,
autosave=True))
# Before the selection is updated, save edits
self.state.register_callback(self.state.pre_selection_updated, self.save_edits)
# When the selection is updated, refresh the nav selection and textbox
self.state.register_callback(self.state.selection_updated, self.fix_selection)
self.state.register_callback(self.state.selection_updated, self.update_nav_tree_selected)
self.state.register_callback(self.state.selection_updated, self.update_chapter_nav_tree_selected)
self.state.register_callback(self.state.selection_updated, self.refresh_textbox)
self.state.register_callback(self.state.selection_updated, self.refresh_alt_textbox)
self.state.register_callback(self.state.selection_updated, self.refresh_vis_selection)
self.state.register_callback(self.state.selection_updated, self.refresh_counterfactual_meta)
self.state.register_callback(self.state.selection_updated, self.refresh_display)
self.state.register_callback(self.state.selection_updated, self.modules_selection_updated)
def bind(self, tk_key, f):
def in_edit():
return self.display.mode in ["Edit", "Child Edit"] \
or (self.display.mode == "Visualize" and self.display.vis.textbox) \
or self.has_focus(self.display.search_box) or self.module_textbox_has_focus() \
or self.state.preferences['editable'] and self.has_focus(self.display.textbox)
valid_keys_outside_edit = ["Control", "Alt", "Escape", "Delete", "Command"]
inside_edit = any(v in tk_key for v in valid_keys_outside_edit)
self.root.bind(tk_key, no_junk_args(f if inside_edit else gated_call(f, lambda: not in_edit())))
def setup_key_bindings(self):
attrs = [getattr(self, f) for f in dir(self)]
funcs_with_keys = [f for f in attrs if callable(f) and hasattr(f, "meta") and "keys" in f.meta]
for f in funcs_with_keys:
for tk_key in f.meta["keys"]:
self.bind(tk_key, f)
# Numbers to select children
# TODO fix this
for i in range(1, 6):
i = i % 10
f = lambda _i=i: self.child(idx=_i-1)
self.bind(f"<Key-{i}>", f)
def setup_custom_key_bindings(self, **kwargs):
for tag, properties in self.state.tags.items():
if properties['toggle_key'] != 'None':
f = lambda _tag=tag: self.toggle_tag(_tag)
tk_key = tkinter_keybindings(properties['toggle_key'])
self.bind(tk_key, f)
def build_menus(self):
# Tuple of 4 things: Name, Hotkey display text, tkinter key to bind to, function to call (without arguments)
menu_list = {
"View": [
('Toggle side pane', 'Alt-P', None, no_junk_args(self.toggle_side)),
('Toggle bottom pane', 'Alt-B', None, no_junk_args(self.toggle_bottom)),
('Toggle visualize mode', 'J', None, no_junk_args(self.toggle_visualization_mode)),
('Toggle children', 'Alt-C', None, no_junk_args(self.toggle_show_children)),
"-",
('Reset zoom', 'Ctrl-0', None, no_junk_args(self.reset_zoom)),
('Center view', 'L, Ctrl-L', None, no_junk_args(self.center_view)),
"-",
('Hoist subtree', 'Alt-H', None, no_junk_args(self.hoist)),
('Unhoist subtree', 'Alt-Shift-H', None, no_junk_args(self.unhoist)),
('Unhoist all', '', None, no_junk_args(self.unhoist_all)),
('Collapse node', 'Ctrl-?', None, no_junk_args(self.collapse_node)),
('Collapse subtree', 'Ctrl-minus', None, no_junk_args(self.collapse_subtree)),
('Collapse all except subtree', 'Ctrl-:', None, no_junk_args(self.collapse_all_except_subtree)),
('Expand children', 'Ctrl-\"', None, no_junk_args(self.expand_node)),
('Expand subtree', 'Ctrl-+', None, no_junk_args(self.expand_subtree)),
('Unzip', '', None, no_junk_args(self.unzip_node)),
('Zip chain', '', None, no_junk_args(self.zip_chain)),
('Zip all', '', None, no_junk_args(self.zip_all_chains)),
('Unzip all', '', None, no_junk_args(self.unzip_all)),
('Show hidden children', '', None, no_junk_args(self.show_hidden_children)),
('Hide invisible children', '', None, no_junk_args(self.hide_invisible_children)),
],
"Edit": [
('Edit node', 'Ctrl+E', None, no_junk_args(self.toggle_edit_mode)),
('Toggle textbox editable', 'Alt+Shift+E', None, no_junk_args(self.toggle_editable)),
"-",
("New root child", 'Ctrl+Shift+H', None, no_junk_args(self.create_root_child)),
("Create parent", 'Alt-Left', None, no_junk_args(self.create_parent)),
("Change parent", 'Shift-P', None, no_junk_args(self.change_parent)),
("New child", 'H, Ctrl+H, Alt+Right', None, no_junk_args(self.create_child)),
("New sibling", 'Alt+Down', None, no_junk_args(self.create_sibling)),
("Merge with parent", 'Shift+Left', None, no_junk_args(self.merge_parent)),
("Merge with children", 'Shift+Right', None, no_junk_args(self.merge_children)),
("Move up", 'Shift+Up', None, no_junk_args(self.move_up)),
("Move down", 'Shift+Down', None, no_junk_args(self.move_down)),
('Prepend newline', 'N, Ctrl+N', None, no_junk_args(self.prepend_newline)),
('Prepend space', 'Ctrl+Space', None, no_junk_args(self.prepend_space)),
('Copy', 'Ctrl+C', None, no_junk_args(self.copy_text)),
('Delete', 'Backspace', None, no_junk_args(self.delete_node)),
('Delete and reassign children', '', None, no_junk_args(self.delete_node_reassign_children)),
],
"Navigate": [
('Return to root', 'R', None, no_junk_args(self.return_to_root)),
('Save checkpoint', 'Ctrl+T', None, no_junk_args(self.save_checkpoint)),
('Go to checkpoint', 'T', None, no_junk_args(self.goto_checkpoint)),
("Next Bookmark", "D", None, no_junk_args(self.next_bookmark)),
("Prev Bookmark", "A", None, no_junk_args(self.prev_bookmark)),
("Stochastic walk", "W", None, no_junk_args(self.walk)),
("Search ancestry", "Ctrl+F", None, no_junk_args(self.search_ancestry)),
("Search tree", "Ctrl+Shift+F", None, no_junk_args(self.search)),
("Goto node by id", "Ctrl+Shift+G", None, no_junk_args(self.goto_node_dialog)),
],
"Generation": [
('Generation settings', 'Ctrl+shift+p', None, no_junk_args(self.generation_settings_dialog)),
('Generate', 'G, Ctrl+G', None, no_junk_args(self.generate)),
#('View summaries', '', None, no_junk_args(self.view_summaries)),
],
"Memory": [
('AI Memory', 'Ctrl+Shift+M', None, no_junk_args(self.ai_memory)),
('Create memory', 'Ctrl+M', None, no_junk_args(self.add_memory)),
('Node memory', 'Alt+M', None, no_junk_args(self.node_memory)),
],
"Flags": [
("Tag node", None, None, no_junk_args(self.tag_node)),
("Add tag", None, None, no_junk_args(self.add_tag)),
("Configure tags", None, None, no_junk_args(self.configure_tags)),
("Mark node visited", None, None, lambda: self.set_visited(True)),
("Mark node unvisited", None, None, lambda: self.set_visited(False)),
("Edit chapter", "Ctrl+Y", None, no_junk_args(self.chapter_dialog)),
],
"Settings": [
('Preferences', 'Ctrl+P', None, no_junk_args(self.preferences)),
('Generation settings', 'Ctrl+shift+p', None, no_junk_args(self.generation_settings_dialog)),
('Inline generation settings', None, None, no_junk_args(self.inline_generation_settings_dialog)),
('Visualization settings', 'Ctrl+U', None, no_junk_args(self.visualization_settings_dialog)),
('Chat settings', None, None, no_junk_args(self.chat_dialog)),
('Model config', None, None, no_junk_args(self.model_config_dialog)),
#('Settings', None, None, no_junk_args(self.settings))
],
"Developer": [
('Run code', 'Ctrl+Shift+B', None, no_junk_args(self.run)),
],
"Info": [
("Tree statistics", "I", None, no_junk_args(self.info_dialog)),
('Multimedia', 'U', None, no_junk_args(self.multimedia_dialog)),
('Node metadata', 'Ctrl+Shift+N', None, no_junk_args(self.node_info_dialogue)),
],
}
return menu_list
@metadata(name='Debug', keys=["<Control-Shift-KeyPress-D>"])
def debug(self, event=None):
#self.display.textbox.fix_selection()
#self.write_textbox_changes()
# makes text copyable
self.display.textbox.bind("<Button>", lambda e: self.display.textbox.focus_set())
#print(self.display.textbox)
print('debug')
# @metadata(name='Generate', keys=["<Control-G>", "<Control-KeyPress-G>"])
# def generate(self, event=None):
# self.generate_dialog()
@metadata(name='View summaries', keys=["<Control-KeyPress-V>"])
def view_summaries(self, event=None):
self.view_summaries_dialog()
@metadata(name='AI Memory', keys=["<Control-KeyPress-M>"])
def ai_memory(self, event=None):
self.ai_memory_dialog()
@metadata(name='Create memory', keys=["<Control-KeyPress-M>"])
def add_memory(self, event=None):
self.add_memory_dialog()
@metadata(name='Node memory', keys=["<Alt-KeyPress-M>"])
def node_memory(self, event=None):
self.node_memory_dialog()
@metadata(name='Search ancestry', keys=["<Control-KeyPress-F>"])
def search_ancestry(self, event=None):
self.search_ancestry_dialog()
@metadata(name='Search tree', keys=["<Control-Shift-KeyPress-F>"])
def search(self, event=None):
self.search_dialog()
#################################
# Navigation
#################################
# @metadata(name=, keys=, display_key=)
@metadata(name="Next", keys=["<period>", "<Return>", "<Control-period>"], display_key=">")
def next(self, node=None):
node = node if node else self.state.selected_node
self.select_node(node=self.state.node(self.state.find_next(node=node, visible_filter=self.in_nav)))
@metadata(name="Prev", keys=["<comma>", "<Control-comma>"], display_key="<",)
def prev(self, node=None):
node = node if node else self.state.selected_node
self.select_node(node=self.state.node(self.state.find_prev(node=node, visible_filter=self.in_nav)))
@metadata(name="Go to parent", keys=["<Left>", "<Control-Left>"], display_key="←")
def parent(self, node=None):
node = node if node else self.state.selected_node
parent = self.state.parent(node)
if parent:
self.select_node(node=parent)
@metadata(name="Go to child", keys=["<Right>", "<Control-Right>"], display_key="→")
def child(self, node=None, idx=0):
node = node if node else self.state.selected_node
child = self.state.child(node, idx, filter=self.in_nav)
if child is not None:
self.select_node(node=child)
@metadata(name="Go to next sibling", keys=["<Down>", "<Control-Down>"], display_key="↓")
def next_sibling(self, node=None):
node = node if node else self.state.selected_node
self.select_node(node=self.state.sibling(node, 1, filter=self.in_nav))
@metadata(name="Go to previous Sibling", keys=["<Up>", "<Control-Up>"], display_key="↑")
def prev_sibling(self, node=None):
node = node if node else self.state.selected_node
self.select_node(node=self.state.sibling(node, -1, filter=self.in_nav))
@metadata(name="Walk", keys=["<Key-w>", "<Control-w>"], display_key="w")
def walk(self, node=None, filter=None):
# TODO custom probs
node = node if node else self.state.selected_node
filter = filter if filter else self.in_nav
if 'children' in node and len(node['children']) > 0:
chosen_child = stochastic_transition(node, mode='leaves', filter=filter)
self.select_node(node=chosen_child)
@metadata(name="Return to root", keys=["<Key-r>", "<Control-r>"], display_key="r")
def return_to_root(self):
self.select_node(node=self.state.root())
@metadata(name="Save checkpoint", keys=["<Control-t>"], display_key="ctrl-t")
def save_checkpoint(self, node=None):
node = node if node else self.state.selected_node
self.state.checkpoint = node['id']
self.state.tree_updated(edit=[node['id']])
@metadata(name="Go to checkpoint", keys=["<Key-t>"], display_key="t")
def goto_checkpoint(self):
if self.state.checkpoint:
self.select_node(node=self.state.node(self.state.checkpoint))
@metadata(name="Nav Select")
def nav_select(self, *, node_id, open=False):
if not node_id or node_id == self.state.selected_node_id:
return
if self.change_parent.meta["click_mode"]:
self.change_parent(node=self.state.node(node_id))
# TODO This causes infinite recursion from the vis node. Need to change how updating open status works
# Update the open state of the node based on the nav bar
# node = self.state.node(node_id]
# node["open"] = self.display.nav_tree.item(node["id"], "open")
#self.state.select_node(node_id)
self.select_node(node=self.state.node(node_id), open=open)
# figure out scope and whether should add, edit, delete
@metadata(name="Tag")
def toggle_tag(self, tag, node=None):
node = node if node else self.state.selected_node
self.state.toggle_tag(node, tag)
self.state.update_tree_tag_changed(node, tag)
@metadata(name="Toggle prompt", keys=["<asterisk>"], display_key="")
def toggle_prompt(self, node=None):
self.state.preferences['show_prompt'] = not self.state.preferences['show_prompt']
self.refresh_textbox()
def next_tag(self, tag, node=None):
node = node if node else self.state.selected_node
next_tag_id = self.state.find_next(node=node, filter=lambda node: self.state.has_tag_attribute(node, tag),
visible_filter=self.in_nav)
self.select_node(self.state.node(next_tag_id))
def prev_tag(self, tag, node=None):
node = node if node else self.state.selected_node
prev_tag_id = self.state.find_prev(node=node, filter=lambda node: self.state.has_tag_attribute(node, tag),
visible_filter=self.in_nav)
self.select_node(self.state.node(prev_tag_id))
@metadata(name="Go to next bookmark", keys=["<Key-d>", "<Control-d>"])
def next_bookmark(self):
self.next_tag(self.state.preferences.get("nav_tag", "bookmark"))
@metadata(name="Go to prev bookmark", keys=["<Key-a>", "<Control-a>"])
def prev_bookmark(self):
self.prev_tag(self.state.preferences.get("nav_tag", "bookmark"))
@metadata(name="Center view", keys=["<Key-l>", "<Control-l>"])
def center_view(self):
#self.display.vis.center_view_on_canvas_coords(*self.display.vis.node_coords[self.state.selected_node_id])
#self.display.vis.center_view_on_node(self.state.selected_node)
pass
def update_read_color(self, old_node, node):
if self.display.mode == 'Read':
pass
# nca_node, index = nearest_common_ancestor(old_node, node, self.state.tree_node_dict)
# ancestor_indices = self.state.ancestor_text_indices(self.state.selected_node)
# ancestor_end_indices = [ind[1] for ind in ancestor_indices]
# nca_end_index = ancestor_end_indices[index]
# self.display.textbox.tag_delete("old")
# self.display.textbox.tag_add("old",
# "1.0",
# f"1.0 + {nca_end_index} chars")
# self.display.textbox.tag_config("old", foreground=history_color())
# self.display.textbox.see(f"1.0 + {nca_end_index} chars")
# pass
#print('coloring text')
@metadata(name="In nav")
def in_nav(self, node):
return self.display.nav_tree.exists(node['id'])
@metadata(name="Select node")
def select_node(self, node, noscroll=False, ask_reveal=True, open=True):
if node == self.state.selected_node:
return
if not self.in_nav(node):
if ask_reveal:
if not self.ask_reveal(node):
return
else:
self.reveal_node(node)
if not self.state.selected_node:
# if no selected node (probably previous node was deleted), just select
self.state.select_node(node['id'])
#print('writing textbox changes')
self.write_textbox_changes()
if open:
node['open'] = True
self.refresh_nav_node(node)
else:
node['open'] = self.display.nav_tree.item(node["id"], "open")
self.nav_history.append(self.state.selected_node_id)
self.undo_history = []
self.state.select_node(node['id'])
# if self.state.preferences['coloring'] == 'read':
# old_node = self.state.selected_node
# self.state.select_node(node['id'], noscroll=True)
# if old_node:
# self.update_read_color(old_node, node)
# else:
#self.state.select_node(node['id'])
@metadata(name="Update text")
def update_text(self, text, node=None):
node = node if node else self.state.selected_node
self.state.update_text(node, text, save_revision_history=self.state.preferences['revision_history'])
@metadata(name="<", keys=["<Command-minus>", "<Alt-minus>"])
def prev_selection(self):
if self.nav_history:
self.undo_history.append(self.state.selected_node_id)
self.state.select_node(self.nav_history.pop())
@metadata(name=">", keys=["<Command-equal>", "<Alt-equal>"])
def next_selection(self):
if self.undo_history:
self.nav_history.append(self.state.selected_node_id)
self.state.select_node(self.undo_history.pop())
@metadata(name="Undo")
def undo_action(self):
# this navigates to the parent of the last non-AI generated node
ancestry = self.state.ancestry(self.state.selected_node)
if len(ancestry) > 1:
for ancestor in ancestry[::-1]:
if not self.state.is_AI_generated(ancestor):
self.parent(node=ancestor)
return
@metadata(name="Rewind")
def rewind(self, node=None, filter=None):
# navigates to previous node in ancestry with more than one child
node = node if node else self.state.selected_node
ancestry = self.state.ancestry(node)[:-1]
if len(ancestry) > 1:
for ancestor in ancestry[::-1]:
children = self.get_children(ancestor, filter=filter)
if len(children) > 1:
self.select_node(ancestor)
return
@metadata(name="Reroll")
def alternate(self, node=None):
node = node if node else self.state.selected_node
next_sibling = self.state.sibling(node, wrap=True, filter=self.in_nav)
self.select_node(next_sibling)
#################################
# Getters
#################################
@metadata(name="Get children")
def get_children(self, node=None, filter=None):
node = node if node else self.state.selected_node
filter = filter if filter else self.in_nav
return filtered_children(node, filter)
@metadata(name="Hidden children")
def get_hidden_children(self, node=None):
node = node if node else self.state.selected_node
return [n for n in node['children'] if not self.state.visible(n)]
@metadata(name="Text")
def get_text(self, node_id=None, raw=False):
node_id = node_id if node_id else self.state.selected_node_id
return self.state.text(self.state.node(node_id), raw=raw)#self.state.node(node_id)['text']
@metadata(name="Get floating notes")
def get_floating_notes(self, tag='note', node=None):
node = node if node else self.state.selected_node
ancestry = self.state.ancestry(node)
notes = []
for ancestor in reversed(ancestry):
for child in ancestor['children']:
if self.state.has_tag(child, tag) and child != node:
notes.append(child)
return notes
@metadata(name="Pinned")
def get_pinned(self):
pinned = self.state.tagged_nodes(tag='pinned')
return pinned
@metadata(name="Prompt")
def prompt(self, node=None):
node = node if node else self.state.selected_node
return self.state.prompt(node)
@metadata(name="Get text attribute")
def get_text_attribute(self, attribute, node=None):
node = node if node else self.state.selected_node
return self.state.get_text_attribute(node, attribute)
#################################
# Node operations
#################################
@metadata(name="New root child", keys=["<Control-Shift-KeyPress-H>"], display_key="ctrl-shift-h" )
def create_root_child(self):
self.create_child(node=self.state.root())
@metadata(name="New Child", keys=["<h>", "<Control-h>", "<Command-Right>", "<Alt-Right>"], display_key="h",)
def create_child(self, node=None, update_selection=True, toggle_edit=True):
node = node if node else self.state.selected_node
new_child = self.state.create_child(parent=node)
self.state.tree_updated(add=[new_child['id']])
self.state.node_creation_metadata(new_child, source='prompt')
if update_selection:
self.select_node(new_child, ask_reveal=False)
if self.display.mode == "Read" and toggle_edit and not self.state.preferences['editable']:
self.toggle_edit_mode()
return new_child
@metadata(name="New Sibling", keys=["<Command-Down>", "<Alt-Down>"], display_key="Command-down")
def create_sibling(self, node=None, toggle_edit=True):
node = node if node else self.state.selected_node
new_sibling = self.state.create_sibling(node=node)
self.state.tree_updated(add=[new_sibling['id']])
self.select_node(new_sibling, ask_reveal=False)
self.state.node_creation_metadata(new_sibling)
if self.display.mode == "Read" and toggle_edit and not self.state.preferences['editable']:
self.toggle_edit_mode()
return new_sibling
@metadata(name="New Parent", keys=["<Command-Left>", "<Alt-Left>"], display_key="Command-left")
def create_parent(self, node=None):
node = node if node else self.state.selected_node
new_parent = self.state.create_parent(node=node)
self.state.tree_updated(add=[new_parent['id']])
self.state.tree_updated(add=[n['id'] for n in subtree_list(new_parent, filter=self.in_nav)])
self.state.node_creation_metadata(new_parent, source='prompt')
return new_parent
@metadata(name="New note")
def new_note(self, node=None):
node = node if node else self.state.selected_node
new_child = self.state.create_child(parent=node)
self.state.tag_node(new_child, 'note')
if self.state.visible(new_child):
self.state.tree_updated(add=[new_child['id']])
else:
self.state.tree_updated()
return new_child
@metadata(name="Change Parent", keys=["<Shift-P>"], display_key="shift-p", selected_node=None, click_mode=False)
def change_parent(self, node=None, click_mode=False):
node = node if node else self.state.selected_node
if self.change_parent.meta["selected_node"] is None:
self.display.change_cursor("fleur")
self.change_parent.meta["selected_node"] = node
self.change_parent.meta["click_mode"] = click_mode
else:
self.display.change_cursor("arrow")
self.state.change_parent(node=self.change_parent.meta["selected_node"], new_parent_id=node["id"])
self.state.tree_updated(add=[n['id'] for n in subtree_list(self.change_parent.meta["selected_node"],
filter=self.in_nav)])
self.change_parent.meta["selected_node"] = None
self.change_parent.meta["click_mode"] = False
@metadata(name="Merge with Parent", keys=["<Shift-Left>"], display_key="shift-left",)
def merge_parent(self, node=None):
node = node if node else self.state.selected_node
if not self.state.is_mutable(node):
self.immutable_popup(node)
return
parent = self.state.parent(node)
if not self.state.is_mutable(parent):
self.immutable_popup(parent)
return
self.state.merge_with_parent(node=node)
self.state.tree_updated(add=[n['id'] for n in subtree_list(parent, filter=self.in_nav)])
self.select_node(parent, ask_reveal=False)
@metadata(name="Merge with children", keys=["<Shift-Right>"], display_key="shift-right")
def merge_children(self, node=None):
node = node if node else self.state.selected_node
if not node['children']:
print('no children')
return
if not self.state.is_mutable(node):
self.immutable_popup(node)
return
children = node['children']
for child in children:
if not self.state.is_mutable(child):
self.immutable_popup(child)
return
visible_children = filtered_children(node, filter=self.in_nav)
self.state.merge_with_children(node=node)
self.state.tree_updated(add=[n['id'] for n in subtree_list(self.state.parent(node), filter=self.in_nav)])
if visible_children:
self.select_node(visible_children[0])
@metadata(name="Move up", keys=["<Shift-Up>"], display_key="shift-up")
def move_up(self, node=None):
if node is None:
node = self.state.selected_node
self.state.shift(node, -1)
self.state.tree_updated(add=[n['id'] for n in subtree_list(self.state.parent(node), filter=self.in_nav)])
@metadata(name="Move down", keys=["<Shift-Down>"], display_key="shift-down")
def move_down(self, node=None):
if node is None:
node = self.state.selected_node
self.state.shift(node, 1)
self.state.tree_updated(add=[n['id'] for n in subtree_list(self.state.parent(node), filter=self.in_nav)])
@metadata(name="Generate", keys=["<g>", "<Control-g>"], display_key="g")
def generate(self, node=None, **kwargs):
if node is None:
node = self.state.selected_node
try:
node["open"] = True
self.display.nav_tree.item(node['id'], open=True)
except Exception as e:
print(str(e))
self.state.generate_continuations(node=node, **kwargs)
@metadata(name="Retry")
def retry(self, node=None):
# if node has a next sibling, select it
# otherwise, navigate to parent, generate again, and select first newly generated node
node = node if node else self.state.selected_node
if self.state.is_AI_generated(node):
next_sibling = self.state.sibling(node, wrap=False, filter=self.in_nav)
self.select_node(next_sibling)
if not self.state.is_AI_generated(self.state.selected_node):
# moved to parent
self.generate(update_selection=True, placeholder="")
else:
self.generate(update_selection=True, placeholder="")
# def propagate_wavefunction(self):
# if self.display.mode == "Multiverse":
# if self.display.multiverse.active_wavefunction():
# active_node = self.display.multiverse.active_info()
# start_position = (active_node['x'], active_node['y'])
# multiverse, ground_truth, prompt = self.state.generate_greedy_multiverse(max_depth=4, prompt=active_node['prefix'],
# unnormalized_amplitude=active_node['amplitude'],
# ground_truth="",
# threshold=0.04,
# engine='ada')
# else:
# start_position = (0, 0)
# multiverse, ground_truth, prompt = self.state.generate_greedy_multiverse(max_depth=4, ground_truth="",
# threshold=0.04,
# engine='ada')
# self.display.multiverse.draw_multiverse(multiverse=multiverse, ground_truth=ground_truth,
# start_position=start_position, prompt=prompt)
@metadata(name="Delete", keys=["<BackSpace>", "<Control-BackSpace>"], display_key="«")
def delete_node(self, node=None, reassign_children=False, ask=True, ask_text="Delete node and subtree?", refresh_nav=True):
node = node if node else self.state.selected_node
if not node:
return
if node == self.state.root():
messagebox.showerror("Root", "Cannot delete root")
return
if ask:
result = messagebox.askquestion("Delete", ask_text, icon='warning')
if result != 'yes':
return False
next_sibling = self.state.sibling(node, wrap=False, filter=self.in_nav)
self.state.delete_node(node=node, reassign_children=reassign_children)
if self.state.selected_node_id == node['id']:
self.select_node(next_sibling)
if self.in_nav(node) and refresh_nav:
self.state.tree_updated(delete=[node['id']])
else:
self.state.tree_updated()
return True
@metadata(name="Delete children")
def delete_children(self, ask=True, node=None):
node = node if node else self.state.selected_node
if not node:
return
if not node['children']:
return
children = node['children']
child_ids = [n['id'] for n in children]
if ask:
result = messagebox.askquestion("Delete", f"Delete {len(children)} children and subtrees?", icon='warning')
if result != 'yes':
return
for child in children:
self.delete_node(child, reassign_children=False, ask=False, refresh_nav=False)
self.state.tree_updated(delete=child_ids)#n['id'] for n in subtree_list(node, filter=self.in_nav) if n != node])
@metadata(name="Archive children")
def archive_children(self, node=None):
node = node if node else self.state.selected_node
if not node:
return
if not node['children']:
return
children = node['children']
for child in children:
self.state.tag_node(child, "archived")
self.state.update_tree_tag_changed(node, "archived")
@metadata(name="Delete and reassign children")
def delete_node_reassign_children(self, node=None):
self.delete_node(node=node, reassign_children=True)
@metadata(name="Enter text", keys=["<Control-bar>"], display_key="Ctrl-Return")
def enter_text(self):
self.display.vis.delete_textbox()
@metadata(name="Escape textbox", keys=["<Escape>"], display_key="Escape")
def escape(self):
self.display.vis.delete_textbox(save=False)
if self.display.in_edit_mode:
self.display.set_mode("Read")
self.refresh_textbox()
if self.change_parent.meta["selected_node"]:
self.display.change_cursor("arrow")
self.change_parent.meta["selected_node"] = None
@metadata(name="Edit history", keys=[], display_key="", persistent_id=None)
def edit_history(self, index):
if self.display.mode == "Read":
_, selected_ancestor = self.index_to_ancestor(index)
self.state.select_node(selected_ancestor["id"])
self.toggle_edit_mode()
@metadata(name="Goto history", keys=[], display_key="")
def goto_history(self, index):
if self.display.mode == "Read":
_, selected_ancestor = self.index_to_ancestor(index)
self.nav_select(node_id=selected_ancestor["id"])
@metadata(name="Split node", keys=[], display_key="")
def split_node(self, index, change_selection=True, node=None):
node = node if node else self.state.selected_node
if self.display.mode == "Read":
if self.state.is_template(node):
#TODO
return
self.write_textbox_changes()
ancestor_index, selected_ancestor = self.index_to_ancestor(index)
ancestor_end_indices = [ind[1] for ind in self.state.ancestor_text_indices(node)]
negative_offset = ancestor_end_indices[ancestor_index] - index
split_index = len(selected_ancestor['text']) - negative_offset
new_parent, _ = self.state.split_node(selected_ancestor, split_index)
self.state.tree_updated(add=[new_parent['id']])
self.state.tree_updated(add=[n['id'] for n in subtree_list(self.state.parent(new_parent), filter=self.in_nav)])
if change_selection:
self.nav_select(node_id=new_parent["id"])
# TODO deal with metadata
def zip_chain(self, node=None):
node = node if node else self.state.selected_node
self.state.zip_chain(node, filter=self.in_nav, refresh_nav=True, update_selection=True)
@metadata(name="Unzip node")
def unzip_node(self, node=None):
node = node if node else self.state.selected_node
self.state.unzip(node, filter=self.state.visible)
def zip_all_chains(self):
editable = False
# temporarily disable editable textbox so that text isn't overwritten
if self.state.preferences.get('editable', False):
editable = True
self.state.update_user_frame(update={'preferences': {'editable': False}})
self.state.zip_all_chains(filter=self.in_nav)
self.state.tree_updated(rebuild=True, write=False)
self.state.select_node(self.state.tree_raw_data['root']['id'], write=False)
# TODO hack
if editable:
self.state.update_user_frame(update={'preferences': {'editable': True}})
def unzip_all(self):
editable = False
if self.state.preferences.get('editable', False):
editable = True
self.state.update_user_frame(update={'preferences': {'editable': False}})
self.state.unzip_all(filter=self.state.visible)
self.state.tree_updated(rebuild=True, write=False)
self.state.select_node(self.state.tree_raw_data['root']['id'], write=False)
# TODO hack
if editable:
self.state.update_user_frame(update={'preferences': {'editable': True}})
#################################
# Textbox
#################################
@metadata(name="Textbox menu")
def textbox_menu(self, char_index, tk_current, e):
_, clicked_node = self.index_to_ancestor(char_index)
node_range = self.node_range(clicked_node)
self.display.textbox.tag_add("node_select", f"1.0 + {node_range[0]} chars", f"1.0 + {node_range[1]} chars")
menu = tk.Menu(self.display.vis.textbox, tearoff=0)
if self.display.textbox.tag_ranges("sel"):
menu.add_command(label="Copy", command=lambda: self.display.textbox.copy_selected())
menu.add_command(label="Go", command=lambda: self.select_node(clicked_node))
menu.add_command(label="Edit", command=lambda: self.edit_in_module(clicked_node))
menu.add_command(label="Split", command=lambda: self.split_node(char_index, change_selection=True))
# TODO
#menu.add_command(label="Generate")
#menu.add_command(label="Add memory")
# splice_menu = tk.Menu(menu, tearoff=0)
# splice_menu.add_command(label="Obliviate")
# splice_menu.add_command(label="Inject")
# if self.display.textbox.tag_ranges("sel"):
# splice_menu.add_command(label="Mask")
# splice_menu.add_command(label="Open window from...")
# splice_menu.add_command(label="Open window to...")
# splice_menu.add_command(label="Custom splice")
#menu.add_cascade(label="Splice", menu=splice_menu)
#menu.add_command(label="Annotate")
# if there is text selected
if self.display.textbox.tag_ranges("sel"):
selection_menu = tk.Menu(menu, tearoff=0)
#TODO
# save_menu = tk.Menu(selection_menu, tearoff=0)
# save_menu.add_command(label="Save text")
# save_menu.add_command(label="Save window")
# save_menu.add_command(label="Save as memory")
# selection_menu.add_cascade(label="Save", menu=save_menu)
#selection_menu.add_command(label="Substitute")
transform_menu = tk.Menu(menu, tearoff=0)
transform_menu.add_command(label="Prose to script", command=lambda: self.open_selection_in_transformer(template='./config/transformers/prose_to_script.json'))
selection_menu.add_cascade(label="Transform", menu=transform_menu)
#selection_menu.add_command(label="Transform", command=lambda: self.open_selection_in_transformer())
# selection_menu.add_command(label="Link")
# selection_menu.add_command(label="Add summary")
# selection_menu.add_command(label="Delete")
menu.add_cascade(label="Selection", menu=selection_menu)
tag_menu = tk.Menu(menu, tearoff=0)
# tag_menu.add_command(label="Pin")
# tag_menu.add_command(label="Tag...")
# menu.add_cascade(label="Tag", menu=tag_menu)
#menu.add_command(label="Copy node")
menu.add_command(label="Copy id", command=lambda: self.copy_id(clicked_node))
menu.tk_popup(e.x_root, e.y_root + 8)
def generate_template(self, inputs, template_file):
# open template file json
with open(template_file, 'r') as f:
template_dict = json.load(f)
template = template_dict['template']
prompt = eval(f'f"""{template}"""')
def refresh_textbox(self, **kwargs):
#print('refresh textbox')
if not self.state.tree_raw_data or not self.state.selected_node:
return
self.display.clear_selection_tags(self.display.textbox)
self.display.textbox.configure(font=Font(family="Georgia", size=self.state.preferences['font_size']),
spacing1=self.state.preferences['paragraph_spacing'],
spacing2=self.state.preferences['line_spacing'],
background=edit_color() if self.state.preferences["editable"] or self.display.mode == "Edit" else bg_color())
#self.display.textbox.tag_config("node_select", font=Font(family="Georgia", size=self.state.preferences['font_size'], weight="bold"))
# Fill textbox with text history, disable editing
if self.display.mode == "Read":
#self.display.textbox.tag_config("sel", background="black", foreground=text_color())
self.display.textbox.configure(state="normal")
self.display.textbox.delete("1.0", "end")
# if self.state.preferences.get('show_prompt', False):
# self.display.textbox.insert("end-1c", self.state.prompt(self.state.selected_node))
# else:
if self.state.preferences['coloring'] in ('edit', 'read'):
self.display.textbox.tag_config('ooc_history', foreground=ooc_color())
self.display.textbox.tag_config('history', foreground=history_color())
else:
self.display.textbox.tag_config('ooc_history', foreground=text_color())
self.display.textbox.tag_config('history', foreground=text_color())
ancestry = self.state.ancestor_text_list(self.state.selected_node)
#self.ancestor_end_indices = indices
history = ''
for node_text in ancestry[:-1]:
# "end" includes the automatically inserted new line
history += node_text
selected_text = self.state.text(self.state.selected_node)#self.state.selected_node["text"]
prompt_length = self.state.generation_settings['prompt_length'] - len(selected_text)
in_context = history[-prompt_length:]
if prompt_length < len(history):
out_context = history[:len(history)-prompt_length]
self.display.textbox.insert("end-1c", out_context, "ooc_history")
self.display.textbox.insert("end-1c", in_context, "history")
history_end = self.display.textbox.index(tk.END)
self.display.textbox.insert("end-1c", selected_text)
active_append_text = self.state.get_text_attribute(self.state.selected_node, 'active_append')
if active_append_text:
self.display.textbox.insert("end-1c", active_append_text)
self.tag_prompts()
if not kwargs.get('noscroll', False):
self.display.textbox.update_idletasks()
if self.state.preferences['coloring'] == 'edit':
self.display.textbox.see(tk.END)
else:
self.display.textbox.see(history_end)
if not self.state.preferences.get('editable', False):
self.display.textbox.configure(state="disabled")
# makes text copyable
#self.display.textbox.bind("<Button>", lambda event: self.display.textbox.focus_set())
# Textbox to edit mode, fill with single node
elif self.display.mode == "Edit":
self.display.textbox.configure(state="normal")
self.display.textbox.delete("1.0", "end")
# TODO depending on show template mode
self.display.textbox.insert("1.0", self.state.selected_node["text"])#self.state.text(self.state.selected_node))#self.state.selected_node["text"])
self.display.textbox.see(tk.END)
# self.display.secondary_textbox.delete("1.0", "end")
# self.display.secondary_textbox.insert("1.0", self.state.selected_node.get("active_text", ""))
self.display.textbox.focus()
# makes text copyable
#self.display.textbox.bind("<Button>", lambda event: self.display.textbox.focus_set())