Added ability to send documents with different prompts based off of d…

…ocument types

utils/constants.py

@@ -47,7 +47,6 @@
     ollama_model TEXT,
     CONSTRAINT docs_fk
     FOREIGN KEY (doc_id) REFERENCES documents (doc_id)
-
 )"""
 CREATE_SCHEDULE_TABLE="""CREATE TABLE IF NOT EXISTS schedule(
     job_id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -66,6 +65,15 @@
     date_updated DATE
 )
 """
+CREATE_DOCUMENT_TYPE_LINK_TABLE="""CREATE TABLE IF NOT EXISTS document_type_link(
+       type_id INTEGER,
+       doc_id INTEGER,
+       CONSTRAINT docs_fk
+       FOREIGN KEY (doc_id) REFERENCES documents (doc_id)
+       CONSTRAINT docs_type_fk
+       FOREIGN KEY (type_id) REFERENCES document_type (type_id)
+)
+"""
 CREATE_PROMPT_TABLE="""CREATE TABLE IF NOT EXISTS prompts(
     prompt_id INTEGER PRIMARY KEY AUTOINCREMENT,
     name TEXT,
@@ -77,10 +85,9 @@
 )
 """
 
-INSTRUCTION_PROMPT="""Do not greet or describe what you are thinking when responding. Stay strictly to the format below. Only summarize using the guidelines below."""
+INSTRUCTION_PROMPT="""Do not greet or describe what you are thinking when responding. Stay strictly to the format below."""
 
-DEFAULT_SYSTEM_PROMPT ="""Race: [<Year> Name of Race]
-Summary: [Summarize the event that occurred in the document. Where applicable, include the driver(s), team(s), or organization(s) involved. Where applicable, include any penalties and what regulation was breached as part of the penalty.]
-[doc_data]""".lstrip().rstrip()
+DEFAULT_SYSTEM_PROMPT ="""Summarize the event that occurred in the document. Where applicable, include the driver(s), team(s), or organization(s) involved. Where applicable, include any penalties and what regulation was breached as part of the penalty.
+Raw Data: [doc_data]""".lstrip().rstrip()
 
 BASE_FIA_URL = "https://www.fia.com"

utils/db.py

@@ -761,6 +761,39 @@ def get_prompt_by_link_id(conn, link_id):
         logging.error(f"Error retrieving prompt by link_id {link_id}: {e}")
         return False
 
+def insert_document_type_link(conn, type_id, doc_id):
+    try:
+        cursor = conn.cursor()
+        query = "INSERT INTO document_type_link (type_id, doc_id) VALUES (?, ?)"
+        cursor.execute(query, (type_id, doc_id))
+        conn.commit()
+        logging.info(f"Inserted into document_type_link: type_id={type_id}, doc_id={doc_id}")
+    except sqlite3.Error as e:
+        conn.rollback()
+        logging.error(f"Error inserting into document_type_link: {e}")
+
+def get_document_type_links(conn, type_id=None, doc_id=None):
+    try:
+        cursor = conn.cursor()
+        if type_id and doc_id:
+            query = "SELECT * FROM document_type_link WHERE type_id = ? AND doc_id = ?"
+            cursor.execute(query, (type_id, doc_id))
+        elif type_id:
+            query = "SELECT * FROM document_type_link WHERE type_id = ?"
+            cursor.execute(query, (type_id,))
+        elif doc_id:
+            query = "SELECT * FROM document_type_link WHERE doc_id = ?"
+            cursor.execute(query, (doc_id,))
+        else:
+            query = "SELECT * FROM document_type_link"
+            cursor.execute(query)
+
+        result = cursor.fetchall()
+        return result
+    except sqlite3.Error as e:
+        logging.error(f"Error retrieving from document_type_link: {e}")
+        return []
+
 #CHATGPT GENERATED CODE END#
 
 def get_config_obj(conn):
@@ -792,6 +825,7 @@ def get_conn():
     cur.execute(constants.CREATE_SCHEDULE_TABLE)
     cur.execute(constants.CREATE_DOCUMENT_TYPE_TABLE)
     cur.execute(constants.CREATE_PROMPT_TABLE)
+    cur.execute(constants.CREATE_DOCUMENT_TYPE_LINK_TABLE)
 
     if first_run:
         insert_config(conn, 'dev', False, True, '', '')

utils/utils.py

@@ -180,11 +180,13 @@ def summarize_data(prompt):
         logging.error(exc_type, fname, exc_tb.tb_lineno, e)
         return ""
 
-def build_prompt(pdf_data):
+def build_prompt(pdf_data, doc_type):
     conn = db.get_conn()
-    sys_prompt = db.get_system_prompt(conn)[2]
-
-    prompt = f"{constants.INSTRUCTION_PROMPT}\n{get_fun_prompt()}\n{sys_prompt}".replace("[doc_data]", pdf_data)
+    custom_prompt = db.get_prompt_by_link_id(conn, doc_type[0])
+    prompt = db.get_system_prompt(conn)[2]
+    if custom_prompt:
+        prompt = custom_prompt
+    prompt = f"{constants.INSTRUCTION_PROMPT}\n{get_fun_prompt()}\n{prompt}".replace("[doc_data]", pdf_data)
     print(prompt)
     return prompt
 
@@ -240,8 +242,10 @@ def date_string_time(date):
 def send_document(send_id):
     conn = db.get_conn()
     send_row = db.join_document_send_documents_webhooks(conn, send_id)[0]
+    print(send_row)
     webhook_url = send_row["webhooks"][0]["webhook_link"]
     title = send_row["document_name"]
+    doc_type = get_doc_type_from_doc_title(title)
     doc_url = send_row["document_link"]
     doc_id = send_row["doc_id"]
     doc_time = datetime.fromisoformat(send_row["document_date"]).astimezone(pytz.utc)
@@ -251,7 +255,7 @@ def send_document(send_id):
     # if doc_summary is None or doc_summary[3] == "":
     file_path = get_file_from_url(doc_url)
     pdf_data = get_pdf_data(file_path)
-    prompt = build_prompt(pdf_data)
+    prompt = build_prompt(pdf_data, doc_type)
     summary = f"**Document Date:**\n{date_string(doc_time)} ({date_string_time(doc_time_est)})\n\n{summarize_data(prompt)}"
     ollama_url = db.get_config_ollama_url(conn)
     ollama_model = db.get_config_ollama_model(conn)
@@ -266,6 +270,16 @@ def send_document(send_id):
     # emit('sent_document', {"id":send_id, "status": status})
     return status
 
+def get_doc_type_from_doc_title(title):
+    conn = db.get_conn()
+    allowed_doc_types = db.get_all_document_types(conn)
+    for dts in allowed_doc_types:
+        enum = dts[1].lower()
+        if enum in title.lower():
+            return dts
+    return None
+
+
 def queue_document(send_id):
     conn = db.get_conn()
     queue_row = db.join_document_send_documents_webhooks(conn, send_id)[0]

web/app.py

@@ -91,7 +91,7 @@ def config_ollama():
         system_prompt = db.get_system_prompt(conn)
         webhook_prompts = db.get_prompts_by_type(conn, "WEBHOOK")
         doctype_prompts = db.get_prompts_by_type(conn, "DOCTYPE")
-        doctypes = db.get_active_document_types(conn)
+        doctypes = db.get_all_document_types(conn)
         print(doctypes)
         if system_prompt is None:
             db.insert_prompt(conn, "DEFAULT_SYSTEM", constants.DEFAULT_SYSTEM_PROMPT, "SYSTEM", None)