Skip to content

Commit a206d6a

Browse files
hgohelMelleKoning
authored andcommitted
ChatWithTree addon
1 parent 57891d1 commit a206d6a

File tree

7 files changed

+1228
-1
lines changed

7 files changed

+1228
-1
lines changed

ChatWithTree/ChatWithTree.gpr.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
# ------------------------------------------------------------------------
2+
#
3+
# Register the Gramplet ChatWithTree
4+
#
5+
# ------------------------------------------------------------------------
6+
register(
7+
GRAMPLET,
8+
id="ChatWithTree", # Unique ID for your addon
9+
name=_("Chat With Tree Interactive Addon"), # Display name in Gramps, translatable
10+
description=_("Chat With Tree with the help of AI Large Language Model, needs litellm module"),
11+
version = '0.0.21',
12+
gramps_target_version="6.0", # Specify the Gramps version you are targeting
13+
status=EXPERIMENTAL,
14+
audience = DEVELOPER,
15+
fname="ChatWithTree.py", # The main Python file for your Gramplet
16+
# The 'gramplet' argument points to the class name in your main file
17+
gramplet="ChatWithTreeClass",
18+
gramplet_title=_("Chat With Tree"),
19+
authors = ["Melle Koning"],
20+
authors_email = ["[email protected]"],
21+
height=18,
22+
# addon needs litellm python module
23+
requires_mod=['litellm'],
24+
)

ChatWithTree/ChatWithTree.py

Lines changed: 342 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,342 @@
1+
#
2+
# Gramps - a GTK+/GNOME based genealogy program
3+
#
4+
# Copyright (C) 2025 Melle Koning
5+
#
6+
# This program is free software; you can redistribute it and/or modify
7+
# it under the terms of the GNU General Public License as published by
8+
# the Free Software Foundation; either version 2 of the License, or
9+
# (at your option) any later version.
10+
#
11+
# This program is distributed in the hope that it will be useful,
12+
# but WITHOUT ANY WARRANTY; without even the implied warranty of
13+
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14+
# GNU General Public License for more details.
15+
#
16+
# You should have received a copy of the GNU General Public License
17+
# along with this program; if not, write to the Free Software
18+
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
19+
#
20+
# ChatWithTree.py
21+
import logging
22+
LOG = logging.getLogger(".")
23+
LOG.debug("loading chatwithtree")
24+
# ==============================================================================
25+
# Standard Python libraries
26+
# ==============================================================================
27+
import gi
28+
gi.require_version("Gtk", "3.0")
29+
from gi.repository import Gtk, Gdk
30+
from gi.repository import GLib
31+
32+
# ==============================================================================
33+
# GRAMPS API
34+
# ==============================================================================
35+
from gramps.gen.plug import Gramplet
36+
from gramps.gen.const import GRAMPS_LOCALE as glocale
37+
_ = glocale.get_addon_translator(__file__).gettext
38+
39+
from chatwithllm import IChatLogic, ChatWithLLM, YieldType
40+
41+
try:
42+
from ChatWithTreeBot import ChatBot
43+
except ImportError as e:
44+
LOG.warning(e)
45+
raise ImportError("Failed to import ChatBot from chatbot module: " + str(e))
46+
47+
LOG.debug("ChatWithTree file header loaded successfully.")
48+
49+
ONE_SECOND = 1000 # milliseconds
50+
51+
# ==============================================================================
52+
# Gramplet Class Definition
53+
# ==============================================================================
54+
class ChatWithTreeClass(Gramplet):
55+
"""
56+
A simple interactive Gramplet that takes user input and provides a reply.
57+
58+
This version uses a Gtk.ListBox to create a dynamic, chat-like interface
59+
with styled message "balloons" for user input and system replies.
60+
"""
61+
62+
def __init__(self, parent=None, **kwargs):
63+
"""
64+
The constructor for the Gramplet.
65+
We call the base class constructor here. The GUI is built in the
66+
init() method.
67+
"""
68+
# Call the base class constructor. This is a mandatory step.
69+
Gramplet.__init__(self, parent, **kwargs)
70+
71+
def init(self):
72+
"""
73+
This method is called by the Gramps framework after the Gramplet
74+
has been fully initialized. We build our GUI here.
75+
"""
76+
# Build our custom GUI widgets.
77+
self.vbox = self._build_gui()
78+
79+
# The Gramplet's container widget is found via `self.gui`.
80+
# We first remove the default textview...
81+
self.gui.get_container_widget().remove(self.gui.textview)
82+
# ... and then we add our new vertical box.
83+
self.gui.get_container_widget().add(self.vbox)
84+
85+
# Show all widgets.
86+
self.vbox.show()
87+
# db change signal
88+
self.dbstate.connect('database-changed', self.change_db)
89+
90+
# Instantiate the chat logic class. This decouples the logic from the UI.
91+
# Choose ChatWIthLLM for simple reverse chat
92+
# self.chat_logic = ChatWithLLM()
93+
# Choose Chatbot for chat with Tree
94+
self.chat_logic = None
95+
#self.chat_logic = ChatBot(self)
96+
97+
def change_db(self, db):
98+
"""
99+
This method is called when the database is opened or closed.
100+
The 'dbstate' parameter is the current database state object.
101+
"""
102+
# Add the initial message to the list box.
103+
self._add_message_row(_("Database change detected"), YieldType.PARTIAL)
104+
105+
if self.dbstate.db:
106+
LOG.debug("Database handle is now available. Initializing chatbot.")
107+
# The database is open, so it is now safe to instantiate the chatbot
108+
# and pass the Gramplet instance with a valid db handle.
109+
self.chat_logic = ChatBot(self)
110+
else:
111+
LOG.debug("Database is closed. Chatbot logic is reset.")
112+
self.chat_logic = None
113+
114+
def _build_gui(self):
115+
"""
116+
Creates all the GTK widgets for the Gramplet's user interface.
117+
Returns the top-level container widget.
118+
"""
119+
# Create the main vertical box to hold all our widgets.
120+
vbox = Gtk.Box(orientation=Gtk.Orientation.VERTICAL, spacing=6)
121+
122+
# -------------------
123+
# 1. Chat History Section
124+
# -------------------
125+
# We use a Gtk.ListBox to hold our chat "balloons".
126+
self.chat_listbox = Gtk.ListBox()
127+
# Set a name for CSS styling.
128+
self.chat_listbox.set_name("chat-listbox")
129+
# Ensure the listbox is a single-column list.
130+
self.chat_listbox.set_selection_mode(Gtk.SelectionMode.NONE)
131+
132+
# We need a reference to the scrolled window to control its scrolling.
133+
self.scrolled_window = Gtk.ScrolledWindow()
134+
self.scrolled_window.set_hexpand(True)
135+
self.scrolled_window.set_vexpand(True)
136+
self.scrolled_window.add(self.chat_listbox)
137+
vbox.pack_start(self.scrolled_window, True, True, 0)
138+
139+
# Apply CSS styling for the chat balloons.
140+
self._apply_css_styles()
141+
142+
# -------------------
143+
# 2. Input Section
144+
# -------------------
145+
input_hbox = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL, spacing=6)
146+
147+
self.input_entry = Gtk.Entry()
148+
self.input_entry.set_placeholder_text(_("Type a message..."))
149+
self.input_entry.connect("activate", self.on_process_button_clicked)
150+
input_hbox.pack_start(self.input_entry, True, True, 0)
151+
152+
self.process_button = Gtk.Button(label=_("Send"))
153+
self.process_button.connect("clicked", self.on_process_button_clicked)
154+
input_hbox.pack_start(self.process_button, False, False, 0)
155+
156+
vbox.pack_start(input_hbox, False, False, 0)
157+
158+
# Add the initial message to the list box.
159+
self._add_message_row(_("Chat with Tree initialized. Type /help for help."), YieldType.PARTIAL)
160+
161+
return vbox
162+
163+
def _apply_css_styles(self):
164+
"""
165+
Defines and applies CSS styles to the Gramplet's widgets.
166+
"""
167+
css_provider = Gtk.CssProvider()
168+
css = """
169+
#chat-listbox {
170+
background-color: white;
171+
}
172+
.message-box {
173+
background-color: #f0f0f0; /* Default background */
174+
padding: 10px;
175+
margin: 5px;
176+
border-radius: 15px;
177+
}
178+
.user-message-box {
179+
background-color: #dcf8c6; /* Light green for user messages */
180+
}
181+
.tree-reply-box {
182+
background-color: #d1e2f4; /* Light blue for replies */
183+
}
184+
.tree-toolcall-box {
185+
background-color: #fce8b2; /* Light yellow for tool calls */
186+
}
187+
"""
188+
css_provider.load_from_data(css.encode('utf-8'))
189+
screen = Gdk.Screen.get_default()
190+
context = Gtk.StyleContext()
191+
context.add_provider_for_screen(screen, css_provider, Gtk.STYLE_PROVIDER_PRIORITY_APPLICATION)
192+
193+
# We need to set up a style context on the chat listbox
194+
style_context = self.chat_listbox.get_style_context()
195+
style_context.add_class("message-box") # This won't work on the listbox itself, but it's good practice.
196+
197+
def _add_message_row(self, text:str, reply_type: YieldType):
198+
"""
199+
Creates a new message "balloon" widget and adds it to the listbox.
200+
"""
201+
# Create a horizontal box to act as the message container.
202+
hbox = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL)
203+
hbox.set_spacing(6)
204+
205+
# Create the message "balloon" box.
206+
message_box = Gtk.Box(orientation=Gtk.Orientation.HORIZONTAL)
207+
message_box.get_style_context().add_class("message-box")
208+
209+
# Create the label for the text.
210+
message_label = Gtk.Label(label=text)
211+
message_label.set_halign(Gtk.Align.START)
212+
message_label.set_line_wrap(True)
213+
message_label.set_max_width_chars(80) # Limit width to prevent it from spanning the entire window.
214+
message_box.pack_start(message_label, True, True, 0)
215+
216+
if reply_type == YieldType.USER:
217+
message_box.get_style_context().add_class("user-message-box")
218+
# Align the message balloon to the right.
219+
hbox.set_halign(Gtk.Align.END)
220+
elif reply_type in (YieldType.PARTIAL, YieldType.TOOL_CALL):
221+
message_box.get_style_context().add_class("tree-toolcall-box")
222+
# Align the message balloon to the left.
223+
hbox.set_halign(Gtk.Align.CENTER)
224+
225+
elif reply_type == YieldType.FINAL:
226+
message_box.get_style_context().add_class("tree-reply-box")
227+
# Align the message balloon to the left.
228+
hbox.set_halign(Gtk.Align.START)
229+
230+
# Add the message balloon to the main horizontal container.
231+
hbox.add(message_box)
232+
233+
# Add the whole row to the listbox.
234+
self.chat_listbox.add(hbox)
235+
self.chat_listbox.show_all()
236+
237+
# The goal is to scroll down after adding a row to the box
238+
# after one full second
239+
# so that Gtk has time to redraw the listbox in that time
240+
GLib.timeout_add(ONE_SECOND, self.scroll_to_bottom)
241+
242+
return message_label
243+
244+
def scroll_to_bottom(self):
245+
"""
246+
Helper function to scroll the listbox to the end.
247+
This runs on the main GTK thread after a redraw.
248+
"""
249+
adj = self.scrolled_window.get_vadjustment()
250+
adj.set_value(adj.get_upper())
251+
252+
# Return False to run the callback only once
253+
return GLib.SOURCE_REMOVE
254+
255+
def _get_reply_on_idle(self):
256+
"""
257+
This is a separate method and to be called via GLib.idle_add
258+
Goal: gets the reply from chatbot and updates the UI.
259+
It runs when the main loop is idle, therefore we return
260+
either GLib.SOURCE_CONTINUE in case there are more replies,
261+
or GLib.SOURCE_REMOVE when the iteration is done
262+
"""
263+
try:
264+
265+
# Using a sentinel object to check for exhaustion
266+
SENTINEL = object()
267+
# use the assigned self.reply_iterator iterator to get the next reply
268+
result = next(self.reply_iterator, SENTINEL)
269+
if result is SENTINEL:
270+
# end of iteration, no replies from iterator
271+
return GLib.SOURCE_REMOVE
272+
# unpack the result tuple
273+
reply_type, content = result
274+
if reply_type == YieldType.PARTIAL:
275+
# sometimes there is no content in the partial yield
276+
# if there is, it is usually an explained strategy what the
277+
# model will do to achieve the final result
278+
self._add_message_row(content, reply_type)
279+
if reply_type == YieldType.TOOL_CALL:
280+
if self.current_tool_call_label is None:
281+
self.current_tool_call_label = self._add_message_row(content, reply_type)
282+
else:
283+
# This is a subsequent tool call. Update the existing label.
284+
# We append the new content to the old content for a streaming effect.
285+
existing_text = self.current_tool_call_label.get_text()
286+
self.current_tool_call_label.set_text(existing_text + " " + content)
287+
elif reply_type == YieldType.FINAL:
288+
# Final reply from the chatbot
289+
# We let the iterator SENTINEL take care of returning Glib.SOURCE_REMOVE
290+
self._add_message_row(content, reply_type)
291+
292+
return GLib.SOURCE_CONTINUE
293+
294+
except Exception as e:
295+
# Handle potential errors from the get_reply function
296+
error_message = f"Error: {type(e).__name__} - {e}"
297+
self._add_message_row(f"Type 'help' for help. \n{error_message}", YieldType.PARTIAL)
298+
299+
return GLib.SOURCE_REMOVE # Stop the process on error
300+
301+
# This function must return False to be removed from the idle handler list.
302+
# If it returns True, it will be called again on the next idle loop.
303+
return False
304+
305+
def on_process_button_clicked(self, widget):
306+
"""
307+
Callback function when the 'Send' button is clicked or 'Enter' is pressed.
308+
"""
309+
# Check if the chat_logic instance has been set.
310+
# This handles the case where the addon is loaded for the first time
311+
# on an already running Gramps session.
312+
if self.chat_logic is None:
313+
self._add_message_row(
314+
"The ChatWithTree addon is not yet initialized. Please reload Gramps or select a database.",
315+
YieldType.FINAL
316+
)
317+
return
318+
# Normal handling of user input
319+
user_input = self.input_entry.get_text()
320+
self.input_entry.set_text("")
321+
if user_input.strip():
322+
# Add the user's message to the chat.
323+
self._add_message_row(f"{user_input}", YieldType.USER)
324+
325+
# Now, schedule the reply-getting logic to run when the main loop is idle.
326+
self.reply_iterator = self.chat_logic.get_reply(user_input)
327+
self.current_tool_call_label = None
328+
329+
GLib.idle_add(self._get_reply_on_idle)
330+
331+
332+
def main(self):
333+
"""
334+
This method is called when the Gramplet needs to update its content.
335+
"""
336+
pass
337+
338+
def destroy(self):
339+
"""
340+
Clean up resources when the Gramplet is closed.
341+
"""
342+
Gramplet.destroy(self)

0 commit comments

Comments
 (0)