oceansweep commited on
Commit
28d9ecb
1 Parent(s): 02b34b4

Upload 3 files

Browse files
App_Function_Libraries/DB/Character_Chat_DB.py ADDED
@@ -0,0 +1,684 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # character_chat_db.py
2
+ # Database functions for managing character cards and chat histories.
3
+ # #
4
+ # Imports
5
+ import configparser
6
+ import sqlite3
7
+ import json
8
+ import os
9
+ import sys
10
+ from typing import List, Dict, Optional, Tuple, Any, Union
11
+
12
+ from App_Function_Libraries.Utils.Utils import get_database_dir, get_project_relative_path, get_database_path
13
+ from Tests.Chat_APIs.Chat_APIs_Integration_test import logging
14
+
15
+ #
16
+ #######################################################################################################################
17
+ #
18
+ #
19
+
20
+ def ensure_database_directory():
21
+ os.makedirs(get_database_dir(), exist_ok=True)
22
+
23
+ ensure_database_directory()
24
+
25
+
26
+ # Construct the path to the config file
27
+ config_path = get_project_relative_path('Config_Files/config.txt')
28
+
29
+ # Read the config file
30
+ config = configparser.ConfigParser()
31
+ config.read(config_path)
32
+
33
+ # Get the chat db path from the config, or use the default if not specified
34
+ chat_DB_PATH = config.get('Database', 'chatDB_path', fallback=get_database_path('chatDB.db'))
35
+ print(f"Chat Database path: {chat_DB_PATH}")
36
+
37
+ ########################################################################################################
38
+ #
39
+ # Functions
40
+
41
+ # FIXME - Setup properly and test/add documentation for its existence...
42
+ def initialize_database():
43
+ """Initialize the SQLite database with required tables and FTS5 virtual tables."""
44
+ conn = None
45
+ try:
46
+ conn = sqlite3.connect(chat_DB_PATH)
47
+ cursor = conn.cursor()
48
+
49
+ # Enable foreign key constraints
50
+ cursor.execute("PRAGMA foreign_keys = ON;")
51
+
52
+ # Create CharacterCards table with V2 fields
53
+ cursor.execute("""
54
+ CREATE TABLE IF NOT EXISTS CharacterCards (
55
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
56
+ name TEXT UNIQUE NOT NULL,
57
+ description TEXT,
58
+ personality TEXT,
59
+ scenario TEXT,
60
+ image BLOB,
61
+ post_history_instructions TEXT,
62
+ first_mes TEXT,
63
+ mes_example TEXT,
64
+ creator_notes TEXT,
65
+ system_prompt TEXT,
66
+ alternate_greetings TEXT,
67
+ tags TEXT,
68
+ creator TEXT,
69
+ character_version TEXT,
70
+ extensions TEXT,
71
+ created_at DATETIME DEFAULT CURRENT_TIMESTAMP
72
+ );
73
+ """)
74
+
75
+ # Create CharacterChats table
76
+ cursor.execute("""
77
+ CREATE TABLE IF NOT EXISTS CharacterChats (
78
+ id INTEGER PRIMARY KEY AUTOINCREMENT,
79
+ character_id INTEGER NOT NULL,
80
+ conversation_name TEXT,
81
+ chat_history TEXT,
82
+ is_snapshot BOOLEAN DEFAULT FALSE,
83
+ created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
84
+ FOREIGN KEY (character_id) REFERENCES CharacterCards(id) ON DELETE CASCADE
85
+ );
86
+ """)
87
+
88
+ # Create FTS5 virtual table for CharacterChats
89
+ cursor.execute("""
90
+ CREATE VIRTUAL TABLE IF NOT EXISTS CharacterChats_fts USING fts5(
91
+ conversation_name,
92
+ chat_history,
93
+ content='CharacterChats',
94
+ content_rowid='id'
95
+ );
96
+ """)
97
+
98
+ # Create triggers to keep FTS5 table in sync with CharacterChats
99
+ cursor.executescript("""
100
+ CREATE TRIGGER IF NOT EXISTS CharacterChats_ai AFTER INSERT ON CharacterChats BEGIN
101
+ INSERT INTO CharacterChats_fts(rowid, conversation_name, chat_history)
102
+ VALUES (new.id, new.conversation_name, new.chat_history);
103
+ END;
104
+
105
+ CREATE TRIGGER IF NOT EXISTS CharacterChats_ad AFTER DELETE ON CharacterChats BEGIN
106
+ DELETE FROM CharacterChats_fts WHERE rowid = old.id;
107
+ END;
108
+
109
+ CREATE TRIGGER IF NOT EXISTS CharacterChats_au AFTER UPDATE ON CharacterChats BEGIN
110
+ UPDATE CharacterChats_fts SET conversation_name = new.conversation_name, chat_history = new.chat_history
111
+ WHERE rowid = new.id;
112
+ END;
113
+ """)
114
+
115
+ # Create ChatKeywords table
116
+ cursor.execute("""
117
+ CREATE TABLE IF NOT EXISTS ChatKeywords (
118
+ chat_id INTEGER NOT NULL,
119
+ keyword TEXT NOT NULL,
120
+ FOREIGN KEY (chat_id) REFERENCES CharacterChats(id) ON DELETE CASCADE
121
+ );
122
+ """)
123
+
124
+ # Create indexes for faster searches
125
+ cursor.execute("""
126
+ CREATE INDEX IF NOT EXISTS idx_chatkeywords_keyword ON ChatKeywords(keyword);
127
+ """)
128
+ cursor.execute("""
129
+ CREATE INDEX IF NOT EXISTS idx_chatkeywords_chat_id ON ChatKeywords(chat_id);
130
+ """)
131
+
132
+ conn.commit()
133
+ logging.info("Database initialized successfully.")
134
+ except sqlite3.Error as e:
135
+ logging.error(f"SQLite error occurred during database initialization: {e}")
136
+ if conn:
137
+ conn.rollback()
138
+ raise
139
+ except Exception as e:
140
+ logging.error(f"Unexpected error occurred during database initialization: {e}")
141
+ if conn:
142
+ conn.rollback()
143
+ raise
144
+ finally:
145
+ if conn:
146
+ conn.close()
147
+
148
+ # Call initialize_database() at the start of your application
149
+ def setup_chat_database():
150
+ try:
151
+ initialize_database()
152
+ except Exception as e:
153
+ logging.critical(f"Failed to initialize database: {e}")
154
+ sys.exit(1)
155
+
156
+ setup_chat_database()
157
+
158
+ ########################################################################################################
159
+ #
160
+ # Character Card handling
161
+
162
+ def parse_character_card(card_data: Dict[str, Any]) -> Dict[str, Any]:
163
+ """Parse and validate a character card according to V2 specification."""
164
+ v2_data = {
165
+ 'name': card_data.get('name', ''),
166
+ 'description': card_data.get('description', ''),
167
+ 'personality': card_data.get('personality', ''),
168
+ 'scenario': card_data.get('scenario', ''),
169
+ 'first_mes': card_data.get('first_mes', ''),
170
+ 'mes_example': card_data.get('mes_example', ''),
171
+ 'creator_notes': card_data.get('creator_notes', ''),
172
+ 'system_prompt': card_data.get('system_prompt', ''),
173
+ 'post_history_instructions': card_data.get('post_history_instructions', ''),
174
+ 'alternate_greetings': json.dumps(card_data.get('alternate_greetings', [])),
175
+ 'tags': json.dumps(card_data.get('tags', [])),
176
+ 'creator': card_data.get('creator', ''),
177
+ 'character_version': card_data.get('character_version', ''),
178
+ 'extensions': json.dumps(card_data.get('extensions', {}))
179
+ }
180
+
181
+ # Handle 'image' separately as it might be binary data
182
+ if 'image' in card_data:
183
+ v2_data['image'] = card_data['image']
184
+
185
+ return v2_data
186
+
187
+
188
+ def add_character_card(card_data: Dict[str, Any]) -> Optional[int]:
189
+ """Add or update a character card in the database."""
190
+ conn = sqlite3.connect(chat_DB_PATH)
191
+ cursor = conn.cursor()
192
+ try:
193
+ parsed_card = parse_character_card(card_data)
194
+
195
+ # Check if character already exists
196
+ cursor.execute("SELECT id FROM CharacterCards WHERE name = ?", (parsed_card['name'],))
197
+ row = cursor.fetchone()
198
+
199
+ if row:
200
+ # Update existing character
201
+ character_id = row[0]
202
+ update_query = """
203
+ UPDATE CharacterCards
204
+ SET description = ?, personality = ?, scenario = ?, image = ?,
205
+ post_history_instructions = ?, first_mes = ?, mes_example = ?,
206
+ creator_notes = ?, system_prompt = ?, alternate_greetings = ?,
207
+ tags = ?, creator = ?, character_version = ?, extensions = ?
208
+ WHERE id = ?
209
+ """
210
+ cursor.execute(update_query, (
211
+ parsed_card['description'], parsed_card['personality'], parsed_card['scenario'],
212
+ parsed_card['image'], parsed_card['post_history_instructions'], parsed_card['first_mes'],
213
+ parsed_card['mes_example'], parsed_card['creator_notes'], parsed_card['system_prompt'],
214
+ parsed_card['alternate_greetings'], parsed_card['tags'], parsed_card['creator'],
215
+ parsed_card['character_version'], parsed_card['extensions'], character_id
216
+ ))
217
+ else:
218
+ # Insert new character
219
+ insert_query = """
220
+ INSERT INTO CharacterCards (name, description, personality, scenario, image,
221
+ post_history_instructions, first_mes, mes_example, creator_notes, system_prompt,
222
+ alternate_greetings, tags, creator, character_version, extensions)
223
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
224
+ """
225
+ cursor.execute(insert_query, (
226
+ parsed_card['name'], parsed_card['description'], parsed_card['personality'],
227
+ parsed_card['scenario'], parsed_card['image'], parsed_card['post_history_instructions'],
228
+ parsed_card['first_mes'], parsed_card['mes_example'], parsed_card['creator_notes'],
229
+ parsed_card['system_prompt'], parsed_card['alternate_greetings'], parsed_card['tags'],
230
+ parsed_card['creator'], parsed_card['character_version'], parsed_card['extensions']
231
+ ))
232
+ character_id = cursor.lastrowid
233
+
234
+ conn.commit()
235
+ return character_id
236
+ except sqlite3.IntegrityError as e:
237
+ logging.error(f"Error adding character card: {e}")
238
+ return None
239
+ except Exception as e:
240
+ logging.error(f"Unexpected error adding character card: {e}")
241
+ return None
242
+ finally:
243
+ conn.close()
244
+
245
+ # def add_character_card(card_data: Dict) -> Optional[int]:
246
+ # """Add or update a character card in the database.
247
+ #
248
+ # Returns the ID of the inserted character or None if failed.
249
+ # """
250
+ # conn = sqlite3.connect(chat_DB_PATH)
251
+ # cursor = conn.cursor()
252
+ # try:
253
+ # # Ensure all required fields are present
254
+ # required_fields = ['name', 'description', 'personality', 'scenario', 'image', 'post_history_instructions', 'first_message']
255
+ # for field in required_fields:
256
+ # if field not in card_data:
257
+ # card_data[field] = '' # Assign empty string if field is missing
258
+ #
259
+ # # Check if character already exists
260
+ # cursor.execute("SELECT id FROM CharacterCards WHERE name = ?", (card_data['name'],))
261
+ # row = cursor.fetchone()
262
+ #
263
+ # if row:
264
+ # # Update existing character
265
+ # character_id = row[0]
266
+ # cursor.execute("""
267
+ # UPDATE CharacterCards
268
+ # SET description = ?, personality = ?, scenario = ?, image = ?, post_history_instructions = ?, first_message = ?
269
+ # WHERE id = ?
270
+ # """, (
271
+ # card_data['description'],
272
+ # card_data['personality'],
273
+ # card_data['scenario'],
274
+ # card_data['image'],
275
+ # card_data['post_history_instructions'],
276
+ # card_data['first_message'],
277
+ # character_id
278
+ # ))
279
+ # else:
280
+ # # Insert new character
281
+ # cursor.execute("""
282
+ # INSERT INTO CharacterCards (name, description, personality, scenario, image, post_history_instructions, first_message)
283
+ # VALUES (?, ?, ?, ?, ?, ?, ?)
284
+ # """, (
285
+ # card_data['name'],
286
+ # card_data['description'],
287
+ # card_data['personality'],
288
+ # card_data['scenario'],
289
+ # card_data['image'],
290
+ # card_data['post_history_instructions'],
291
+ # card_data['first_message']
292
+ # ))
293
+ # character_id = cursor.lastrowid
294
+ #
295
+ # conn.commit()
296
+ # return cursor.lastrowid
297
+ # except sqlite3.IntegrityError as e:
298
+ # logging.error(f"Error adding character card: {e}")
299
+ # return None
300
+ # except Exception as e:
301
+ # logging.error(f"Unexpected error adding character card: {e}")
302
+ # return None
303
+ # finally:
304
+ # conn.close()
305
+
306
+
307
+ def get_character_cards() -> List[Dict]:
308
+ """Retrieve all character cards from the database."""
309
+ logging.debug(f"Fetching characters from DB: {chat_DB_PATH}")
310
+ conn = sqlite3.connect(chat_DB_PATH)
311
+ cursor = conn.cursor()
312
+ cursor.execute("SELECT * FROM CharacterCards")
313
+ rows = cursor.fetchall()
314
+ columns = [description[0] for description in cursor.description]
315
+ conn.close()
316
+ characters = [dict(zip(columns, row)) for row in rows]
317
+ #logging.debug(f"Characters fetched from DB: {characters}")
318
+ return characters
319
+
320
+
321
+ def get_character_card_by_id(character_id: Union[int, Dict[str, Any]]) -> Optional[Dict[str, Any]]:
322
+ """
323
+ Retrieve a single character card by its ID.
324
+
325
+ Args:
326
+ character_id: Can be either an integer ID or a dictionary containing character data.
327
+
328
+ Returns:
329
+ A dictionary containing the character card data, or None if not found.
330
+ """
331
+ conn = sqlite3.connect(chat_DB_PATH)
332
+ cursor = conn.cursor()
333
+ try:
334
+ if isinstance(character_id, dict):
335
+ # If a dictionary is passed, assume it's already a character card
336
+ return character_id
337
+ elif isinstance(character_id, int):
338
+ # If an integer is passed, fetch the character from the database
339
+ cursor.execute("SELECT * FROM CharacterCards WHERE id = ?", (character_id,))
340
+ row = cursor.fetchone()
341
+ if row:
342
+ columns = [description[0] for description in cursor.description]
343
+ return dict(zip(columns, row))
344
+ else:
345
+ logging.warning(f"Invalid type for character_id: {type(character_id)}")
346
+ return None
347
+ except Exception as e:
348
+ logging.error(f"Error in get_character_card_by_id: {e}")
349
+ return None
350
+ finally:
351
+ conn.close()
352
+
353
+
354
+ def update_character_card(character_id: int, card_data: Dict) -> bool:
355
+ """Update an existing character card."""
356
+ conn = sqlite3.connect(chat_DB_PATH)
357
+ cursor = conn.cursor()
358
+ try:
359
+ cursor.execute("""
360
+ UPDATE CharacterCards
361
+ SET name = ?, description = ?, personality = ?, scenario = ?, image = ?, post_history_instructions = ?, first_message = ?
362
+ WHERE id = ?
363
+ """, (
364
+ card_data.get('name'),
365
+ card_data.get('description'),
366
+ card_data.get('personality'),
367
+ card_data.get('scenario'),
368
+ card_data.get('image'),
369
+ card_data.get('post_history_instructions', ''),
370
+ card_data.get('first_message', "Hello! I'm ready to chat."),
371
+ character_id
372
+ ))
373
+ conn.commit()
374
+ return cursor.rowcount > 0
375
+ except sqlite3.IntegrityError as e:
376
+ logging.error(f"Error updating character card: {e}")
377
+ return False
378
+ finally:
379
+ conn.close()
380
+
381
+
382
+ def delete_character_card(character_id: int) -> bool:
383
+ """Delete a character card and its associated chats."""
384
+ conn = sqlite3.connect(chat_DB_PATH)
385
+ cursor = conn.cursor()
386
+ try:
387
+ # Delete associated chats first due to foreign key constraint
388
+ cursor.execute("DELETE FROM CharacterChats WHERE character_id = ?", (character_id,))
389
+ cursor.execute("DELETE FROM CharacterCards WHERE id = ?", (character_id,))
390
+ conn.commit()
391
+ return cursor.rowcount > 0
392
+ except sqlite3.Error as e:
393
+ logging.error(f"Error deleting character card: {e}")
394
+ return False
395
+ finally:
396
+ conn.close()
397
+
398
+
399
+ def add_character_chat(character_id: int, conversation_name: str, chat_history: List[Tuple[str, str]], keywords: Optional[List[str]] = None, is_snapshot: bool = False) -> Optional[int]:
400
+ """
401
+ Add a new chat history for a character, optionally associating keywords.
402
+
403
+ Args:
404
+ character_id (int): The ID of the character.
405
+ conversation_name (str): Name of the conversation.
406
+ chat_history (List[Tuple[str, str]]): List of (user, bot) message tuples.
407
+ keywords (Optional[List[str]]): List of keywords to associate with this chat.
408
+ is_snapshot (bool, optional): Whether this chat is a snapshot.
409
+
410
+ Returns:
411
+ Optional[int]: The ID of the inserted chat or None if failed.
412
+ """
413
+ conn = sqlite3.connect(chat_DB_PATH)
414
+ cursor = conn.cursor()
415
+ try:
416
+ chat_history_json = json.dumps(chat_history)
417
+ cursor.execute("""
418
+ INSERT INTO CharacterChats (character_id, conversation_name, chat_history, is_snapshot)
419
+ VALUES (?, ?, ?, ?)
420
+ """, (
421
+ character_id,
422
+ conversation_name,
423
+ chat_history_json,
424
+ is_snapshot
425
+ ))
426
+ chat_id = cursor.lastrowid
427
+
428
+ if keywords:
429
+ # Insert keywords into ChatKeywords table
430
+ keyword_records = [(chat_id, keyword.strip().lower()) for keyword in keywords]
431
+ cursor.executemany("""
432
+ INSERT INTO ChatKeywords (chat_id, keyword)
433
+ VALUES (?, ?)
434
+ """, keyword_records)
435
+
436
+ conn.commit()
437
+ return chat_id
438
+ except sqlite3.Error as e:
439
+ logging.error(f"Error adding character chat: {e}")
440
+ return None
441
+ finally:
442
+ conn.close()
443
+
444
+
445
+ def get_character_chats(character_id: Optional[int] = None) -> List[Dict]:
446
+ """Retrieve all chats, or chats for a specific character if character_id is provided."""
447
+ conn = sqlite3.connect(chat_DB_PATH)
448
+ cursor = conn.cursor()
449
+ if character_id is not None:
450
+ cursor.execute("SELECT * FROM CharacterChats WHERE character_id = ?", (character_id,))
451
+ else:
452
+ cursor.execute("SELECT * FROM CharacterChats")
453
+ rows = cursor.fetchall()
454
+ columns = [description[0] for description in cursor.description]
455
+ conn.close()
456
+ return [dict(zip(columns, row)) for row in rows]
457
+
458
+
459
+ def get_character_chat_by_id(chat_id: int) -> Optional[Dict]:
460
+ """Retrieve a single chat by its ID."""
461
+ conn = sqlite3.connect(chat_DB_PATH)
462
+ cursor = conn.cursor()
463
+ cursor.execute("SELECT * FROM CharacterChats WHERE id = ?", (chat_id,))
464
+ row = cursor.fetchone()
465
+ conn.close()
466
+ if row:
467
+ columns = [description[0] for description in cursor.description]
468
+ chat = dict(zip(columns, row))
469
+ chat['chat_history'] = json.loads(chat['chat_history'])
470
+ return chat
471
+ return None
472
+
473
+
474
+ def search_character_chats(query: str) -> Tuple[List[Dict], str]:
475
+ """
476
+ Search for character chats using FTS5.
477
+
478
+ Args:
479
+ query (str): The search query.
480
+
481
+ Returns:
482
+ Tuple[List[Dict], str]: A list of matching chats and a status message.
483
+ """
484
+ if not query.strip():
485
+ return [], "Please enter a search query."
486
+
487
+ conn = sqlite3.connect(chat_DB_PATH)
488
+ cursor = conn.cursor()
489
+ try:
490
+ # Use parameterized queries to prevent SQL injection
491
+ cursor.execute("""
492
+ SELECT CharacterChats.id, CharacterChats.conversation_name, CharacterChats.chat_history
493
+ FROM CharacterChats_fts
494
+ JOIN CharacterChats ON CharacterChats_fts.rowid = CharacterChats.id
495
+ WHERE CharacterChats_fts MATCH ?
496
+ ORDER BY rank
497
+ """, (query,))
498
+ rows = cursor.fetchall()
499
+ columns = [description[0] for description in cursor.description]
500
+ results = [dict(zip(columns, row)) for row in rows]
501
+ status_message = f"Found {len(results)} chat(s) matching '{query}'."
502
+ return results, status_message
503
+ except Exception as e:
504
+ logging.error(f"Error searching chats with FTS5: {e}")
505
+ return [], f"Error occurred during search: {e}"
506
+ finally:
507
+ conn.close()
508
+
509
+ def update_character_chat(chat_id: int, chat_history: List[Tuple[str, str]]) -> bool:
510
+ """Update an existing chat history."""
511
+ conn = sqlite3.connect(chat_DB_PATH)
512
+ cursor = conn.cursor()
513
+ try:
514
+ chat_history_json = json.dumps(chat_history)
515
+ cursor.execute("""
516
+ UPDATE CharacterChats
517
+ SET chat_history = ?
518
+ WHERE id = ?
519
+ """, (
520
+ chat_history_json,
521
+ chat_id
522
+ ))
523
+ conn.commit()
524
+ return cursor.rowcount > 0
525
+ except sqlite3.Error as e:
526
+ logging.error(f"Error updating character chat: {e}")
527
+ return False
528
+ finally:
529
+ conn.close()
530
+
531
+
532
+ def delete_character_chat(chat_id: int) -> bool:
533
+ """Delete a specific chat."""
534
+ conn = sqlite3.connect(chat_DB_PATH)
535
+ cursor = conn.cursor()
536
+ try:
537
+ cursor.execute("DELETE FROM CharacterChats WHERE id = ?", (chat_id,))
538
+ conn.commit()
539
+ return cursor.rowcount > 0
540
+ except sqlite3.Error as e:
541
+ logging.error(f"Error deleting character chat: {e}")
542
+ return False
543
+ finally:
544
+ conn.close()
545
+
546
+ def fetch_keywords_for_chats(keywords: List[str]) -> List[int]:
547
+ """
548
+ Fetch chat IDs associated with any of the specified keywords.
549
+
550
+ Args:
551
+ keywords (List[str]): List of keywords to search for.
552
+
553
+ Returns:
554
+ List[int]: List of chat IDs associated with the keywords.
555
+ """
556
+ if not keywords:
557
+ return []
558
+
559
+ conn = sqlite3.connect(chat_DB_PATH)
560
+ cursor = conn.cursor()
561
+ try:
562
+ # Construct the WHERE clause to search for each keyword
563
+ keyword_clauses = " OR ".join(["keyword = ?"] * len(keywords))
564
+ sql_query = f"SELECT DISTINCT chat_id FROM ChatKeywords WHERE {keyword_clauses}"
565
+ cursor.execute(sql_query, keywords)
566
+ rows = cursor.fetchall()
567
+ chat_ids = [row[0] for row in rows]
568
+ return chat_ids
569
+ except Exception as e:
570
+ logging.error(f"Error in fetch_keywords_for_chats: {e}")
571
+ return []
572
+ finally:
573
+ conn.close()
574
+
575
+ def save_chat_history_to_character_db(character_id: int, conversation_name: str, chat_history: List[Tuple[str, str]]) -> Optional[int]:
576
+ """Save chat history to the CharacterChats table.
577
+
578
+ Returns the ID of the inserted chat or None if failed.
579
+ """
580
+ return add_character_chat(character_id, conversation_name, chat_history)
581
+
582
+ def migrate_chat_to_media_db():
583
+ pass
584
+
585
+
586
+ def search_db(query: str, fields: List[str], where_clause: str = "", page: int = 1, results_per_page: int = 5) -> List[Dict[str, Any]]:
587
+ """
588
+ Perform a full-text search on specified fields with optional filtering and pagination.
589
+
590
+ Args:
591
+ query (str): The search query.
592
+ fields (List[str]): List of fields to search in.
593
+ where_clause (str, optional): Additional SQL WHERE clause to filter results.
594
+ page (int, optional): Page number for pagination.
595
+ results_per_page (int, optional): Number of results per page.
596
+
597
+ Returns:
598
+ List[Dict[str, Any]]: List of matching chat records with content and metadata.
599
+ """
600
+ if not query.strip():
601
+ return []
602
+
603
+ conn = sqlite3.connect(chat_DB_PATH)
604
+ cursor = conn.cursor()
605
+ try:
606
+ # Construct the MATCH query for FTS5
607
+ match_query = " AND ".join(fields) + f" MATCH ?"
608
+ # Adjust the query with the fields
609
+ fts_query = f"""
610
+ SELECT CharacterChats.id, CharacterChats.conversation_name, CharacterChats.chat_history
611
+ FROM CharacterChats_fts
612
+ JOIN CharacterChats ON CharacterChats_fts.rowid = CharacterChats.id
613
+ WHERE {match_query}
614
+ """
615
+ if where_clause:
616
+ fts_query += f" AND ({where_clause})"
617
+ fts_query += " ORDER BY rank LIMIT ? OFFSET ?"
618
+ offset = (page - 1) * results_per_page
619
+ cursor.execute(fts_query, (query, results_per_page, offset))
620
+ rows = cursor.fetchall()
621
+ columns = [description[0] for description in cursor.description]
622
+ results = [dict(zip(columns, row)) for row in rows]
623
+ return results
624
+ except Exception as e:
625
+ logging.error(f"Error in search_db: {e}")
626
+ return []
627
+ finally:
628
+ conn.close()
629
+
630
+
631
+ def perform_full_text_search_chat(query: str, relevant_chat_ids: List[int], page: int = 1, results_per_page: int = 5) -> \
632
+ List[Dict[str, Any]]:
633
+ """
634
+ Perform a full-text search within the specified chat IDs using FTS5.
635
+
636
+ Args:
637
+ query (str): The user's query.
638
+ relevant_chat_ids (List[int]): List of chat IDs to search within.
639
+ page (int): Pagination page number.
640
+ results_per_page (int): Number of results per page.
641
+
642
+ Returns:
643
+ List[Dict[str, Any]]: List of search results with content and metadata.
644
+ """
645
+ try:
646
+ # Construct a WHERE clause to limit the search to relevant chat IDs
647
+ where_clause = " OR ".join([f"media_id = {chat_id}" for chat_id in relevant_chat_ids])
648
+ if not where_clause:
649
+ where_clause = "1" # No restriction if no chat IDs
650
+
651
+ # Perform full-text search using FTS5
652
+ fts_results = search_db(query, ["content"], where_clause, page=page, results_per_page=results_per_page)
653
+
654
+ filtered_fts_results = [
655
+ {
656
+ "content": result['content'],
657
+ "metadata": {"media_id": result['id']}
658
+ }
659
+ for result in fts_results
660
+ if result['id'] in relevant_chat_ids
661
+ ]
662
+ return filtered_fts_results
663
+ except Exception as e:
664
+ logging.error(f"Error in perform_full_text_search_chat: {str(e)}")
665
+ return []
666
+
667
+
668
+ def fetch_all_chats() -> List[Dict[str, Any]]:
669
+ """
670
+ Fetch all chat messages from the database.
671
+
672
+ Returns:
673
+ List[Dict[str, Any]]: List of chat messages with relevant metadata.
674
+ """
675
+ try:
676
+ chats = get_character_chats() # Modify this function to retrieve all chats
677
+ return chats
678
+ except Exception as e:
679
+ logging.error(f"Error fetching all chats: {str(e)}")
680
+ return []
681
+
682
+ #
683
+ # End of Character_Chat_DB.py
684
+ #######################################################################################################################
App_Function_Libraries/DB/DB_Manager.py CHANGED
@@ -70,6 +70,14 @@ from App_Function_Libraries.DB.SQLite_DB import (
70
  fetch_paginated_data as sqlite_fetch_paginated_data, get_latest_transcription as sqlite_get_latest_transcription, \
71
  mark_media_as_processed as sqlite_mark_media_as_processed,
72
  )
 
 
 
 
 
 
 
 
73
  #
74
  # Local Imports
75
  from App_Function_Libraries.Utils.Utils import load_comprehensive_config, get_database_path, get_project_relative_path
@@ -510,9 +518,9 @@ def delete_prompt(*args, **kwargs):
510
  # Implement Elasticsearch version
511
  raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
512
 
513
- def search_media_database(query: str) -> List[Tuple[int, str, str]]:
514
  if db_type == 'sqlite':
515
- return sqlite_search_media_database(query)
516
  elif db_type == 'elasticsearch':
517
  # Implement Elasticsearch version when available
518
  raise NotImplementedError("Elasticsearch version of search_media_database not yet implemented")
@@ -770,6 +778,93 @@ def get_conversation_name(*args, **kwargs):
770
  # End of Chat-related Functions
771
  ############################################################################################################
772
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
773
  ############################################################################################################
774
  #
775
  # Trash-related Functions
 
70
  fetch_paginated_data as sqlite_fetch_paginated_data, get_latest_transcription as sqlite_get_latest_transcription, \
71
  mark_media_as_processed as sqlite_mark_media_as_processed,
72
  )
73
+ from App_Function_Libraries.DB.Character_Chat_DB import (
74
+ add_character_card as sqlite_add_character_card, get_character_cards as sqlite_get_character_cards, \
75
+ get_character_card_by_id as sqlite_get_character_card_by_id, update_character_card as sqlite_update_character_card, \
76
+ delete_character_card as sqlite_delete_character_card, add_character_chat as sqlite_add_character_chat, \
77
+ get_character_chats as sqlite_get_character_chats, get_character_chat_by_id as sqlite_get_character_chat_by_id, \
78
+ update_character_chat as sqlite_update_character_chat, delete_character_chat as sqlite_delete_character_chat, \
79
+ migrate_chat_to_media_db as sqlite_migrate_chat_to_media_db,
80
+ )
81
  #
82
  # Local Imports
83
  from App_Function_Libraries.Utils.Utils import load_comprehensive_config, get_database_path, get_project_relative_path
 
518
  # Implement Elasticsearch version
519
  raise NotImplementedError("Elasticsearch version of add_media_with_keywords not yet implemented")
520
 
521
+ def search_media_database(*args, **kwargs):
522
  if db_type == 'sqlite':
523
+ return sqlite_search_media_database(*args, **kwargs)
524
  elif db_type == 'elasticsearch':
525
  # Implement Elasticsearch version when available
526
  raise NotImplementedError("Elasticsearch version of search_media_database not yet implemented")
 
778
  # End of Chat-related Functions
779
  ############################################################################################################
780
 
781
+
782
+ ############################################################################################################
783
+ #
784
+ # Character Chat-related Functions
785
+
786
+ def add_character_card(*args, **kwargs):
787
+ if db_type == 'sqlite':
788
+ return sqlite_add_character_card(*args, **kwargs)
789
+ elif db_type == 'elasticsearch':
790
+ # Implement Elasticsearch version
791
+ raise NotImplementedError("Elasticsearch version of add_character_card not yet implemented")
792
+
793
+ def get_character_cards():
794
+ if db_type == 'sqlite':
795
+ return sqlite_get_character_cards()
796
+ elif db_type == 'elasticsearch':
797
+ # Implement Elasticsearch version
798
+ raise NotImplementedError("Elasticsearch version of get_character_cards not yet implemented")
799
+
800
+ def get_character_card_by_id(*args, **kwargs):
801
+ if db_type == 'sqlite':
802
+ return sqlite_get_character_card_by_id(*args, **kwargs)
803
+ elif db_type == 'elasticsearch':
804
+ # Implement Elasticsearch version
805
+ raise NotImplementedError("Elasticsearch version of get_character_card_by_id not yet implemented")
806
+
807
+ def update_character_card(*args, **kwargs):
808
+ if db_type == 'sqlite':
809
+ return sqlite_update_character_card(*args, **kwargs)
810
+ elif db_type == 'elasticsearch':
811
+ # Implement Elasticsearch version
812
+ raise NotImplementedError("Elasticsearch version of update_character_card not yet implemented")
813
+
814
+ def delete_character_card(*args, **kwargs):
815
+ if db_type == 'sqlite':
816
+ return sqlite_delete_character_card(*args, **kwargs)
817
+ elif db_type == 'elasticsearch':
818
+ # Implement Elasticsearch version
819
+ raise NotImplementedError("Elasticsearch version of delete_character_card not yet implemented")
820
+
821
+ def add_character_chat(*args, **kwargs):
822
+ if db_type == 'sqlite':
823
+ return sqlite_add_character_chat(*args, **kwargs)
824
+ elif db_type == 'elasticsearch':
825
+ # Implement Elasticsearch version
826
+ raise NotImplementedError("Elasticsearch version of add_character_chat not yet implemented")
827
+
828
+ def get_character_chats(*args, **kwargs):
829
+ if db_type == 'sqlite':
830
+ return sqlite_get_character_chats(*args, **kwargs)
831
+ elif db_type == 'elasticsearch':
832
+ # Implement Elasticsearch version
833
+ raise NotImplementedError("Elasticsearch version of get_character_chats not yet implemented")
834
+
835
+ def get_character_chat_by_id(*args, **kwargs):
836
+ if db_type == 'sqlite':
837
+ return sqlite_get_character_chat_by_id(*args, **kwargs)
838
+ elif db_type == 'elasticsearch':
839
+ # Implement Elasticsearch version
840
+ raise NotImplementedError("Elasticsearch version of get_character_chat_by_id not yet implemented")
841
+
842
+ def update_character_chat(*args, **kwargs):
843
+ if db_type == 'sqlite':
844
+ return sqlite_update_character_chat(*args, **kwargs)
845
+ elif db_type == 'elasticsearch':
846
+ # Implement Elasticsearch version
847
+ raise NotImplementedError("Elasticsearch version of update_character_chat not yet implemented")
848
+
849
+ def delete_character_chat(*args, **kwargs):
850
+ if db_type == 'sqlite':
851
+ return sqlite_delete_character_chat(*args, **kwargs)
852
+ elif db_type == 'elasticsearch':
853
+ # Implement Elasticsearch version
854
+ raise NotImplementedError("Elasticsearch version of delete_character_chat not yet implemented")
855
+
856
+ def migrate_chat_to_media_db(*args, **kwargs):
857
+ if db_type == 'sqlite':
858
+ return sqlite_migrate_chat_to_media_db(*args, **kwargs)
859
+ elif db_type == 'elasticsearch':
860
+ # Implement Elasticsearch version
861
+ raise NotImplementedError("Elasticsearch version of migrate_chat_to_media_db not yet implemented")
862
+
863
+ #
864
+ # End of Character Chat-related Functions
865
+ ############################################################################################################
866
+
867
+
868
  ############################################################################################################
869
  #
870
  # Trash-related Functions
App_Function_Libraries/DB/SQLite_DB.py CHANGED
@@ -102,8 +102,8 @@ backup_path = get_project_relative_path(backup_path)
102
  db_path = sqlite_path
103
  backup_dir = backup_path
104
 
105
- print(f"Database path: {db_path}")
106
- print(f"Backup directory: {backup_dir}")
107
  #create_automated_backup(db_path, backup_dir)
108
 
109
  # FIXME - Setup properly and test/add documentation for its existence...
@@ -277,9 +277,9 @@ db = Database()
277
 
278
  # Usage example:
279
  if db.table_exists('DocumentVersions'):
280
- logging.info("DocumentVersions table exists")
281
  else:
282
- logging.error("DocumentVersions table does not exist")
283
 
284
 
285
  # Function to create tables with the new media schema
@@ -467,7 +467,7 @@ def check_media_exists(title: str, url: str) -> Optional[int]:
467
  try:
468
  with db.get_connection() as conn:
469
  cursor = conn.cursor()
470
- query = 'SELECT id FROM Media WHERE title = ? AND url = ?'
471
  cursor.execute(query, (title, url))
472
  result = cursor.fetchone()
473
  logging.debug(f"check_media_exists query: {query}")
@@ -603,57 +603,6 @@ def mark_media_as_processed(database, media_id):
603
  # Keyword-related Functions
604
  #
605
 
606
- # Function to add a keyword
607
- def add_keyword(keyword: str) -> int:
608
- keyword = keyword.strip().lower()
609
- with db.get_connection() as conn:
610
- cursor = conn.cursor()
611
- try:
612
- # Insert into Keywords table
613
- cursor.execute('INSERT OR IGNORE INTO Keywords (keyword) VALUES (?)', (keyword,))
614
-
615
- # Get the keyword_id (whether it was just inserted or already existed)
616
- cursor.execute('SELECT id FROM Keywords WHERE keyword = ?', (keyword,))
617
- keyword_id = cursor.fetchone()[0]
618
-
619
- # Check if the keyword exists in keyword_fts
620
- cursor.execute('SELECT rowid FROM keyword_fts WHERE rowid = ?', (keyword_id,))
621
- if not cursor.fetchone():
622
- # If it doesn't exist in keyword_fts, insert it
623
- cursor.execute('INSERT OR IGNORE INTO keyword_fts (rowid, keyword) VALUES (?, ?)', (keyword_id, keyword))
624
-
625
- logging.info(f"Keyword '{keyword}' added or updated with ID: {keyword_id}")
626
- conn.commit()
627
- return keyword_id
628
- except sqlite3.IntegrityError as e:
629
- logging.error(f"Integrity error adding keyword: {e}")
630
- raise DatabaseError(f"Integrity error adding keyword: {e}")
631
- except sqlite3.Error as e:
632
- logging.error(f"Error adding keyword: {e}")
633
- raise DatabaseError(f"Error adding keyword: {e}")
634
-
635
-
636
-
637
- # Function to delete a keyword
638
- def delete_keyword(keyword: str) -> str:
639
- keyword = keyword.strip().lower()
640
- with db.get_connection() as conn:
641
- cursor = conn.cursor()
642
- try:
643
- cursor.execute('SELECT id FROM Keywords WHERE keyword = ?', (keyword,))
644
- keyword_id = cursor.fetchone()
645
- if keyword_id:
646
- cursor.execute('DELETE FROM Keywords WHERE keyword = ?', (keyword,))
647
- cursor.execute('DELETE FROM keyword_fts WHERE rowid = ?', (keyword_id[0],))
648
- conn.commit()
649
- return f"Keyword '{keyword}' deleted successfully."
650
- else:
651
- return f"Keyword '{keyword}' not found."
652
- except sqlite3.Error as e:
653
- raise DatabaseError(f"Error deleting keyword: {e}")
654
-
655
-
656
-
657
  # Function to add media with keywords
658
  def add_media_with_keywords(url, title, media_type, content, keywords, prompt, summary, transcription_model, author,
659
  ingestion_date):
@@ -828,6 +777,59 @@ def ingest_article_to_db(url, title, author, content, keywords, summary, ingesti
828
  return str(e)
829
 
830
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
831
  def fetch_all_keywords() -> List[str]:
832
  try:
833
  with db.get_connection() as conn:
@@ -998,14 +1000,15 @@ def add_media_version(conn, media_id: int, prompt: str, summary: str) -> None:
998
 
999
 
1000
  # Function to search the database with advanced options, including keyword search and full-text search
1001
- def sqlite_search_db(search_query: str, search_fields: List[str], keywords: str, page: int = 1, results_per_page: int = 10):
 
1002
  if page < 1:
1003
  raise ValueError("Page number must be 1 or greater.")
1004
 
1005
  # Prepare keywords by splitting and trimming
1006
  keywords = [keyword.strip().lower() for keyword in keywords.split(',') if keyword.strip()]
1007
 
1008
- with db.get_connection() as conn:
1009
  cursor = conn.cursor()
1010
  offset = (page - 1) * results_per_page
1011
 
@@ -1042,10 +1045,13 @@ def sqlite_search_db(search_query: str, search_fields: List[str], keywords: str,
1042
  params.extend([results_per_page, offset])
1043
 
1044
  cursor.execute(query, params)
1045
- results = cursor.fetchall()
1046
-
1047
- return results
1048
 
 
 
 
 
 
1049
 
1050
  # Gradio function to handle user input and display results with pagination, with better feedback
1051
  def search_and_display(search_query, search_fields, keywords, page):
@@ -1177,8 +1183,9 @@ def is_valid_date(date_string: str) -> bool:
1177
 
1178
 
1179
 
1180
- def add_media_to_database(url, info_dict, segments, summary, keywords, custom_prompt_input, whisper_model, media_type='video', overwrite=False):
1181
- db = Database()
 
1182
  try:
1183
  with db.get_connection() as conn:
1184
  cursor = conn.cursor()
@@ -1365,12 +1372,14 @@ def schedule_chunking(media_id: int, content: str, media_name: str):
1365
  # Functions to manage prompts DB
1366
 
1367
  def create_prompts_db():
 
1368
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1369
  cursor = conn.cursor()
1370
  cursor.executescript('''
1371
  CREATE TABLE IF NOT EXISTS Prompts (
1372
  id INTEGER PRIMARY KEY AUTOINCREMENT,
1373
  name TEXT NOT NULL UNIQUE,
 
1374
  details TEXT,
1375
  system TEXT,
1376
  user TEXT
@@ -1391,22 +1400,42 @@ def create_prompts_db():
1391
  CREATE INDEX IF NOT EXISTS idx_promptkeywords_keyword_id ON PromptKeywords(keyword_id);
1392
  ''')
1393
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1394
 
1395
  def normalize_keyword(keyword):
1396
  return re.sub(r'\s+', ' ', keyword.strip().lower())
1397
 
1398
 
1399
- def add_prompt(name, details, system, user=None, keywords=None):
1400
- if not name or not system:
1401
- return "Name and system prompt are required."
 
 
 
1402
 
1403
  try:
1404
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1405
  cursor = conn.cursor()
1406
  cursor.execute('''
1407
- INSERT INTO Prompts (name, details, system, user)
1408
- VALUES (?, ?, ?, ?)
1409
- ''', (name, details, system, user))
1410
  prompt_id = cursor.lastrowid
1411
 
1412
  if keywords:
@@ -1428,10 +1457,11 @@ def add_prompt(name, details, system, user=None, keywords=None):
1428
 
1429
 
1430
  def fetch_prompt_details(name):
 
1431
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1432
  cursor = conn.cursor()
1433
  cursor.execute('''
1434
- SELECT p.name, p.details, p.system, p.user, GROUP_CONCAT(k.keyword, ', ') as keywords
1435
  FROM Prompts p
1436
  LEFT JOIN PromptKeywords pk ON p.id = pk.prompt_id
1437
  LEFT JOIN Keywords k ON pk.keyword_id = k.id
@@ -1442,6 +1472,7 @@ def fetch_prompt_details(name):
1442
 
1443
 
1444
  def list_prompts(page=1, per_page=10):
 
1445
  offset = (page - 1) * per_page
1446
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1447
  cursor = conn.cursor()
@@ -1458,6 +1489,7 @@ def list_prompts(page=1, per_page=10):
1458
  # This will not scale. For a large number of prompts, use a more efficient method.
1459
  # FIXME - see above statement.
1460
  def load_preset_prompts():
 
1461
  try:
1462
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1463
  cursor = conn.cursor()
@@ -1469,8 +1501,8 @@ def load_preset_prompts():
1469
  return []
1470
 
1471
 
1472
- def insert_prompt_to_db(title, description, system_prompt, user_prompt, keywords=None):
1473
- return add_prompt(title, description, system_prompt, user_prompt, keywords)
1474
 
1475
 
1476
  def get_prompt_db_connection():
@@ -1479,6 +1511,7 @@ def get_prompt_db_connection():
1479
 
1480
 
1481
  def search_prompts(query):
 
1482
  try:
1483
  with get_prompt_db_connection() as conn:
1484
  cursor = conn.cursor()
@@ -1498,6 +1531,7 @@ def search_prompts(query):
1498
 
1499
 
1500
  def search_prompts_by_keyword(keyword, page=1, per_page=10):
 
1501
  normalized_keyword = normalize_keyword(keyword)
1502
  offset = (page - 1) * per_page
1503
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
@@ -1527,6 +1561,7 @@ def search_prompts_by_keyword(keyword, page=1, per_page=10):
1527
 
1528
 
1529
  def update_prompt_keywords(prompt_name, new_keywords):
 
1530
  try:
1531
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1532
  cursor = conn.cursor()
@@ -1557,40 +1592,43 @@ def update_prompt_keywords(prompt_name, new_keywords):
1557
  return f"Database error: {e}"
1558
 
1559
 
1560
- def add_or_update_prompt(title, description, system_prompt, user_prompt, keywords=None):
 
1561
  if not title:
1562
  return "Error: Title is required."
1563
 
1564
  existing_prompt = fetch_prompt_details(title)
1565
  if existing_prompt:
1566
  # Update existing prompt
1567
- result = update_prompt_in_db(title, description, system_prompt, user_prompt)
1568
  if "successfully" in result:
1569
  # Update keywords if the prompt update was successful
1570
  keyword_result = update_prompt_keywords(title, keywords or [])
1571
  result += f" {keyword_result}"
1572
  else:
1573
  # Insert new prompt
1574
- result = insert_prompt_to_db(title, description, system_prompt, user_prompt, keywords)
1575
 
1576
  return result
1577
 
1578
 
1579
  def load_prompt_details(selected_prompt):
 
1580
  if selected_prompt:
1581
  details = fetch_prompt_details(selected_prompt)
1582
  if details:
1583
- return details[0], details[1], details[2], details[3], details[4] # Include keywords
1584
- return "", "", "", "", ""
1585
 
1586
 
1587
- def update_prompt_in_db(title, description, system_prompt, user_prompt):
 
1588
  try:
1589
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1590
  cursor = conn.cursor()
1591
  cursor.execute(
1592
- "UPDATE Prompts SET details = ?, system = ?, user = ? WHERE name = ?",
1593
- (description, system_prompt, user_prompt, title)
1594
  )
1595
  if cursor.rowcount == 0:
1596
  return "No prompt found with the given title."
@@ -1602,6 +1640,7 @@ def update_prompt_in_db(title, description, system_prompt, user_prompt):
1602
  create_prompts_db()
1603
 
1604
  def delete_prompt(prompt_id):
 
1605
  try:
1606
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1607
  cursor = conn.cursor()
@@ -1671,15 +1710,21 @@ def update_media_content(selected_item, item_mapping, content_input, prompt_inpu
1671
  return f"Error updating content: {str(e)}"
1672
 
1673
 
1674
- def search_media_database(query: str) -> List[Tuple[int, str, str]]:
1675
- try:
1676
- with db.get_connection() as conn:
1677
  cursor = conn.cursor()
1678
  cursor.execute("SELECT id, title, url FROM Media WHERE title LIKE ?", (f'%{query}%',))
1679
- results = cursor.fetchall()
1680
- return results
1681
- except sqlite3.Error as e:
1682
- raise Exception(f"Error searching media database: {e}")
 
 
 
 
 
 
1683
 
1684
  def load_media_content(media_id: int) -> dict:
1685
  try:
@@ -2586,7 +2631,6 @@ def get_paginated_files(page: int = 1, results_per_page: int = 50) -> Tuple[List
2586
  #
2587
  # Functions to manage document versions
2588
 
2589
-
2590
  def create_document_version(media_id: int, content: str) -> int:
2591
  logging.info(f"Attempting to create document version for media_id: {media_id}")
2592
  try:
@@ -2910,22 +2954,23 @@ def update_media_table(db):
2910
  # Add chunking_status column if it doesn't exist
2911
  add_missing_column_if_not_exists(db, 'Media', 'chunking_status', "TEXT DEFAULT 'pending'")
2912
 
2913
- # Vector check FIXME/Delete later
2914
- def alter_media_table(db):
2915
- alter_query = '''
2916
- ALTER TABLE Media ADD COLUMN vector_processing INTEGER DEFAULT 0
2917
- '''
2918
- try:
2919
- db.execute_query(alter_query)
2920
- logging.info("Media table altered successfully to include vector_processing column.")
2921
- except Exception as e:
2922
- logging.error(f"Error altering Media table: {str(e)}")
2923
- # If the column already exists, SQLite will throw an error, which we can safely ignore
2924
- if "duplicate column name" not in str(e).lower():
2925
- raise
2926
-
2927
- # Vector check FIXME/Delete later
2928
- alter_media_table(db)
 
2929
 
2930
  #
2931
  # End of Functions to manage media chunks
 
102
  db_path = sqlite_path
103
  backup_dir = backup_path
104
 
105
+ print(f"Media Database path: {db_path}")
106
+ print(f"Media Backup directory: {backup_dir}")
107
  #create_automated_backup(db_path, backup_dir)
108
 
109
  # FIXME - Setup properly and test/add documentation for its existence...
 
277
 
278
  # Usage example:
279
  if db.table_exists('DocumentVersions'):
280
+ logging.debug("DocumentVersions table exists")
281
  else:
282
+ logging.debug("DocumentVersions table does not exist")
283
 
284
 
285
  # Function to create tables with the new media schema
 
467
  try:
468
  with db.get_connection() as conn:
469
  cursor = conn.cursor()
470
+ query = 'SELECT id FROM Media WHERE title = ? OR url = ?'
471
  cursor.execute(query, (title, url))
472
  result = cursor.fetchone()
473
  logging.debug(f"check_media_exists query: {query}")
 
603
  # Keyword-related Functions
604
  #
605
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
606
  # Function to add media with keywords
607
  def add_media_with_keywords(url, title, media_type, content, keywords, prompt, summary, transcription_model, author,
608
  ingestion_date):
 
777
  return str(e)
778
 
779
 
780
+ # Function to add a keyword
781
+ def add_keyword(keyword: str) -> int:
782
+ if not keyword.strip():
783
+ raise DatabaseError("Keyword cannot be empty")
784
+
785
+ keyword = keyword.strip().lower()
786
+ with db.get_connection() as conn:
787
+ cursor = conn.cursor()
788
+ try:
789
+ # Insert into Keywords table
790
+ cursor.execute('INSERT OR IGNORE INTO Keywords (keyword) VALUES (?)', (keyword,))
791
+
792
+ # Get the keyword_id (whether it was just inserted or already existed)
793
+ cursor.execute('SELECT id FROM Keywords WHERE keyword = ?', (keyword,))
794
+ keyword_id = cursor.fetchone()[0]
795
+
796
+ # Check if the keyword exists in keyword_fts
797
+ cursor.execute('SELECT rowid FROM keyword_fts WHERE rowid = ?', (keyword_id,))
798
+ if not cursor.fetchone():
799
+ # If it doesn't exist in keyword_fts, insert it
800
+ cursor.execute('INSERT OR IGNORE INTO keyword_fts (rowid, keyword) VALUES (?, ?)', (keyword_id, keyword))
801
+
802
+ logging.info(f"Keyword '{keyword}' added or updated with ID: {keyword_id}")
803
+ conn.commit()
804
+ return keyword_id
805
+ except sqlite3.IntegrityError as e:
806
+ logging.error(f"Integrity error adding keyword: {e}")
807
+ raise DatabaseError(f"Integrity error adding keyword: {e}")
808
+ except sqlite3.Error as e:
809
+ logging.error(f"Error adding keyword: {e}")
810
+ raise DatabaseError(f"Error adding keyword: {e}")
811
+
812
+
813
+
814
+ # Function to delete a keyword
815
+ def delete_keyword(keyword: str) -> str:
816
+ keyword = keyword.strip().lower()
817
+ with db.get_connection() as conn:
818
+ cursor = conn.cursor()
819
+ try:
820
+ cursor.execute('SELECT id FROM Keywords WHERE keyword = ?', (keyword,))
821
+ keyword_id = cursor.fetchone()
822
+ if keyword_id:
823
+ cursor.execute('DELETE FROM Keywords WHERE keyword = ?', (keyword,))
824
+ cursor.execute('DELETE FROM keyword_fts WHERE rowid = ?', (keyword_id[0],))
825
+ conn.commit()
826
+ return f"Keyword '{keyword}' deleted successfully."
827
+ else:
828
+ return f"Keyword '{keyword}' not found."
829
+ except sqlite3.Error as e:
830
+ raise DatabaseError(f"Error deleting keyword: {e}")
831
+
832
+
833
  def fetch_all_keywords() -> List[str]:
834
  try:
835
  with db.get_connection() as conn:
 
1000
 
1001
 
1002
  # Function to search the database with advanced options, including keyword search and full-text search
1003
+
1004
+ def sqlite_search_db(search_query: str, search_fields: List[str], keywords: str, page: int = 1, results_per_page: int = 10, connection=None):
1005
  if page < 1:
1006
  raise ValueError("Page number must be 1 or greater.")
1007
 
1008
  # Prepare keywords by splitting and trimming
1009
  keywords = [keyword.strip().lower() for keyword in keywords.split(',') if keyword.strip()]
1010
 
1011
+ def execute_query(conn):
1012
  cursor = conn.cursor()
1013
  offset = (page - 1) * results_per_page
1014
 
 
1045
  params.extend([results_per_page, offset])
1046
 
1047
  cursor.execute(query, params)
1048
+ return cursor.fetchall()
 
 
1049
 
1050
+ if connection:
1051
+ return execute_query(connection)
1052
+ else:
1053
+ with db.get_connection() as conn:
1054
+ return execute_query(conn)
1055
 
1056
  # Gradio function to handle user input and display results with pagination, with better feedback
1057
  def search_and_display(search_query, search_fields, keywords, page):
 
1183
 
1184
 
1185
 
1186
+ def add_media_to_database(url, info_dict, segments, summary, keywords, custom_prompt_input, whisper_model, media_type='video', overwrite=False, db=None):
1187
+ if db is None:
1188
+ db = Database()
1189
  try:
1190
  with db.get_connection() as conn:
1191
  cursor = conn.cursor()
 
1372
  # Functions to manage prompts DB
1373
 
1374
  def create_prompts_db():
1375
+ logging.debug("create_prompts_db: Creating prompts database.")
1376
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1377
  cursor = conn.cursor()
1378
  cursor.executescript('''
1379
  CREATE TABLE IF NOT EXISTS Prompts (
1380
  id INTEGER PRIMARY KEY AUTOINCREMENT,
1381
  name TEXT NOT NULL UNIQUE,
1382
+ author TEXT,
1383
  details TEXT,
1384
  system TEXT,
1385
  user TEXT
 
1400
  CREATE INDEX IF NOT EXISTS idx_promptkeywords_keyword_id ON PromptKeywords(keyword_id);
1401
  ''')
1402
 
1403
+ # FIXME - dirty hack that should be removed later...
1404
+ # Migration function to add the 'author' column to the Prompts table
1405
+ def add_author_column_to_prompts():
1406
+ with sqlite3.connect(get_database_path('prompts.db')) as conn:
1407
+ cursor = conn.cursor()
1408
+ # Check if 'author' column already exists
1409
+ cursor.execute("PRAGMA table_info(Prompts)")
1410
+ columns = [col[1] for col in cursor.fetchall()]
1411
+
1412
+ if 'author' not in columns:
1413
+ # Add the 'author' column
1414
+ cursor.execute('ALTER TABLE Prompts ADD COLUMN author TEXT')
1415
+ print("Author column added to Prompts table.")
1416
+ else:
1417
+ print("Author column already exists in Prompts table.")
1418
+
1419
+ add_author_column_to_prompts()
1420
 
1421
  def normalize_keyword(keyword):
1422
  return re.sub(r'\s+', ' ', keyword.strip().lower())
1423
 
1424
 
1425
+ # FIXME - update calls to this function to use the new args
1426
+ def add_prompt(name, author, details, system=None, user=None, keywords=None):
1427
+ logging.debug(f"add_prompt: Adding prompt with name: {name}, author: {author}, system: {system}, user: {user}, keywords: {keywords}")
1428
+ if not name:
1429
+ logging.error("add_prompt: A name is required.")
1430
+ return "A name is required."
1431
 
1432
  try:
1433
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1434
  cursor = conn.cursor()
1435
  cursor.execute('''
1436
+ INSERT INTO Prompts (name, author, details, system, user)
1437
+ VALUES (?, ?, ?, ?, ?)
1438
+ ''', (name, author, details, system, user))
1439
  prompt_id = cursor.lastrowid
1440
 
1441
  if keywords:
 
1457
 
1458
 
1459
  def fetch_prompt_details(name):
1460
+ logging.debug(f"fetch_prompt_details: Fetching details for prompt: {name}")
1461
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1462
  cursor = conn.cursor()
1463
  cursor.execute('''
1464
+ SELECT p.name, p.author, p.details, p.system, p.user, GROUP_CONCAT(k.keyword, ', ') as keywords
1465
  FROM Prompts p
1466
  LEFT JOIN PromptKeywords pk ON p.id = pk.prompt_id
1467
  LEFT JOIN Keywords k ON pk.keyword_id = k.id
 
1472
 
1473
 
1474
  def list_prompts(page=1, per_page=10):
1475
+ logging.debug(f"list_prompts: Listing prompts for page {page} with {per_page} prompts per page.")
1476
  offset = (page - 1) * per_page
1477
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1478
  cursor = conn.cursor()
 
1489
  # This will not scale. For a large number of prompts, use a more efficient method.
1490
  # FIXME - see above statement.
1491
  def load_preset_prompts():
1492
+ logging.debug("load_preset_prompts: Loading preset prompts.")
1493
  try:
1494
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1495
  cursor = conn.cursor()
 
1501
  return []
1502
 
1503
 
1504
+ def insert_prompt_to_db(title, author, description, system_prompt, user_prompt, keywords=None):
1505
+ return add_prompt(title, author, description, system_prompt, user_prompt, keywords)
1506
 
1507
 
1508
  def get_prompt_db_connection():
 
1511
 
1512
 
1513
  def search_prompts(query):
1514
+ logging.debug(f"search_prompts: Searching prompts with query: {query}")
1515
  try:
1516
  with get_prompt_db_connection() as conn:
1517
  cursor = conn.cursor()
 
1531
 
1532
 
1533
  def search_prompts_by_keyword(keyword, page=1, per_page=10):
1534
+ logging.debug(f"search_prompts_by_keyword: Searching prompts by keyword: {keyword}")
1535
  normalized_keyword = normalize_keyword(keyword)
1536
  offset = (page - 1) * per_page
1537
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
 
1561
 
1562
 
1563
  def update_prompt_keywords(prompt_name, new_keywords):
1564
+ logging.debug(f"update_prompt_keywords: Updating keywords for prompt: {prompt_name}")
1565
  try:
1566
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1567
  cursor = conn.cursor()
 
1592
  return f"Database error: {e}"
1593
 
1594
 
1595
+ def add_or_update_prompt(title, author, description, system_prompt, user_prompt, keywords=None):
1596
+ logging.debug(f"add_or_update_prompt: Adding or updating prompt: {title}")
1597
  if not title:
1598
  return "Error: Title is required."
1599
 
1600
  existing_prompt = fetch_prompt_details(title)
1601
  if existing_prompt:
1602
  # Update existing prompt
1603
+ result = update_prompt_in_db(title, author, description, system_prompt, user_prompt)
1604
  if "successfully" in result:
1605
  # Update keywords if the prompt update was successful
1606
  keyword_result = update_prompt_keywords(title, keywords or [])
1607
  result += f" {keyword_result}"
1608
  else:
1609
  # Insert new prompt
1610
+ result = insert_prompt_to_db(title, author, description, system_prompt, user_prompt, keywords)
1611
 
1612
  return result
1613
 
1614
 
1615
  def load_prompt_details(selected_prompt):
1616
+ logging.debug(f"load_prompt_details: Loading prompt details for {selected_prompt}")
1617
  if selected_prompt:
1618
  details = fetch_prompt_details(selected_prompt)
1619
  if details:
1620
+ return details[0], details[1], details[2], details[3], details[4], details[5]
1621
+ return "", "", "", "", "", ""
1622
 
1623
 
1624
+ def update_prompt_in_db(title, author, description, system_prompt, user_prompt):
1625
+ logging.debug(f"update_prompt_in_db: Updating prompt: {title}")
1626
  try:
1627
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1628
  cursor = conn.cursor()
1629
  cursor.execute(
1630
+ "UPDATE Prompts SET author = ?, details = ?, system = ?, user = ? WHERE name = ?",
1631
+ (author, description, system_prompt, user_prompt, title)
1632
  )
1633
  if cursor.rowcount == 0:
1634
  return "No prompt found with the given title."
 
1640
  create_prompts_db()
1641
 
1642
  def delete_prompt(prompt_id):
1643
+ logging.debug(f"delete_prompt: Deleting prompt with ID: {prompt_id}")
1644
  try:
1645
  with sqlite3.connect(get_database_path('prompts.db')) as conn:
1646
  cursor = conn.cursor()
 
1710
  return f"Error updating content: {str(e)}"
1711
 
1712
 
1713
+ def search_media_database(query: str, connection=None) -> List[Tuple[int, str, str]]:
1714
+ def execute_query(conn):
1715
+ try:
1716
  cursor = conn.cursor()
1717
  cursor.execute("SELECT id, title, url FROM Media WHERE title LIKE ?", (f'%{query}%',))
1718
+ return cursor.fetchall()
1719
+ except sqlite3.Error as e:
1720
+ raise Exception(f"Error searching media database: {e}")
1721
+
1722
+ if connection:
1723
+ return execute_query(connection)
1724
+ else:
1725
+ with db.get_connection() as conn:
1726
+ return execute_query(conn)
1727
+
1728
 
1729
  def load_media_content(media_id: int) -> dict:
1730
  try:
 
2631
  #
2632
  # Functions to manage document versions
2633
 
 
2634
  def create_document_version(media_id: int, content: str) -> int:
2635
  logging.info(f"Attempting to create document version for media_id: {media_id}")
2636
  try:
 
2954
  # Add chunking_status column if it doesn't exist
2955
  add_missing_column_if_not_exists(db, 'Media', 'chunking_status', "TEXT DEFAULT 'pending'")
2956
 
2957
+ # DEADCODE
2958
+ # # Vector check FIXME/Delete later
2959
+ # def alter_media_table(db):
2960
+ # alter_query = '''
2961
+ # ALTER TABLE Media ADD COLUMN vector_processing INTEGER DEFAULT 0
2962
+ # '''
2963
+ # try:
2964
+ # db.execute_query(alter_query)
2965
+ # logging.info("Media table altered successfully to include vector_processing column.")
2966
+ # except Exception as e:
2967
+ # logging.error(f"Error altering Media table: {str(e)}")
2968
+ # # If the column already exists, SQLite will throw an error, which we can safely ignore
2969
+ # if "duplicate column name" not in str(e).lower():
2970
+ # raise
2971
+ #
2972
+ # # Vector check FIXME/Delete later
2973
+ # alter_media_table(db)
2974
 
2975
  #
2976
  # End of Functions to manage media chunks