feat: combine consecutive citations

Author: aydin-utting

common/services/transcription_handler_service.py

@@ -12,6 +12,7 @@
 from common.services.exceptions import InteractionFailedError, TranscriptionFailedError
 from common.services.transcription_services.transcription_manager import TranscriptionServiceManager
 from common.settings import get_settings
+from common.templates.template_utils import combine_consecutive_citations
 from common.types import DialogueEntry, TranscriptionJobMessageData
 
 settings = get_settings()
@@ -76,6 +77,7 @@ async def process_interactive_message(chat_id: UUID) -> None:
                         )
 
                 chat_response = await chatbot.chat(messages=chat_history)
+                chat_response = combine_consecutive_citations(chat_response)
                 chat.assistant_content = chat_response
                 chat.status = JobStatus.COMPLETED
 

common/templates/add_citations.py

@@ -0,0 +1,45 @@
+import re
+
+from common.database.postgres_models import DialogueEntry
+from common.llm.client import FastOrBestLLM, create_default_chatbot
+from common.prompts import get_citations_prompt
+
+
+async def add_citations_to_minute(
+    transcript: list[DialogueEntry],
+    initial_draft: str,
+) -> str:
+    chatbot = create_default_chatbot(FastOrBestLLM.FAST)
+    messages = get_citations_prompt(initial_draft, transcript)
+
+    minute = await chatbot.chat(messages)
+
+    minute = combine_consecutive_citations(minute)
+
+    return minute or ""
+
+
+def combine_consecutive_citations(minute: str) -> str:
+    cluster_pattern = re.compile(r"(\[\d+\])+")
+    matches = cluster_pattern.finditer(minute)
+    for match in matches:
+        citation_cluster = match.group()
+        citation_pattern = re.compile(r"\d+")
+        numbers = [int(n.group()) for n in citation_pattern.finditer(citation_cluster)]
+        numbers.sort()
+        # Extract individual numbers from the cluster
+        groups = []
+        for number in numbers:
+            if len(groups) == 0 or abs(groups[-1][-1] - number) > 2:
+                groups.append([number])
+            else:
+                groups[-1].append(number)
+
+        out = ""
+        for citation_group in groups:
+            if len(citation_group) == 1:
+                out += f"[{citation_group[0]}]"
+            else:
+                out += f"[{citation_group[0]}-{citation_group[-1]}]"
+        minute = minute.replace(citation_cluster, out)
+    return minute

common/templates/template_utils.py

@@ -14,6 +14,7 @@ import { useCallback, useEffect } from 'react'
 
 import { CitationPopoverWrapper } from '@/components/ui/citation-popover-wrapper'
 import { useCitationPopover } from '@/hooks/use-citation-popover'
+import { citationRegex, citationRegexWithSpace } from '@/lib/citationRegex'
 import { Transcription } from '@/lib/client'
 import { cn } from '@/lib/utils'
 import posthog from 'posthog-js'
@@ -58,7 +59,7 @@ function SimpleEditor({
           props: {
             decorations(state) {
               const decorations: Decoration[] = []
-              const citationRegex = /(\s?)\[(\d+)\]/g
+              const citationRegex = citationRegexWithSpace
 
               state.doc.descendants((node, pos) => {
                 if (node.isText) {
@@ -93,7 +94,7 @@ function SimpleEditor({
                 const domNode = event.target as HTMLElement
 
                 if (domNode.classList.contains('citation-link')) {
-                  const match = domNode.textContent?.match(/\[(\d+)\]/)
+                  const match = domNode.textContent?.match(citationRegex)
                   if (match) {
                     const index = parseInt(match[1], 10)
                     const rect = domNode.getBoundingClientRect()

frontend/app/transcriptions/[transcriptionId]/MinuteTab/components/editor/tiptap-editor.tsx

@@ -7,6 +7,7 @@ import { MinuteVersionSelect } from '@/app/transcriptions/[transcriptionId]/Minu
 import { NewMinuteDialog } from '@/app/transcriptions/[transcriptionId]/MinuteTab/NewMinuteDialog'
 import { Button } from '@/components/ui/button'
 import CopyButton from '@/components/ui/copy-button'
+import { citationRegex, citationRegexWithSpace } from '@/lib/citationRegex'
 import {
   MinuteListItem,
   MinuteVersionResponse,
@@ -86,10 +87,10 @@ export function MinuteEditor({
   }, [form, minuteVersion])
   const htmlContent = form.watch('html')
   const contentToCopy = useMemo(() => {
-    return htmlContent?.replaceAll(/\s?\[(\d+)\]/g, '') || ''
+    return htmlContent?.replaceAll(citationRegexWithSpace, '') || ''
   }, [htmlContent])
   const hasCitations = useMemo(() => {
-    return !!htmlContent?.match(/\[(\d+)\]/)
+    return !!htmlContent?.match(citationRegex)
   }, [htmlContent])
   useEffect(() => {}, [htmlContent])
   const { mutate: saveEdit } = useMutation({

frontend/app/transcriptions/[transcriptionId]/MinuteTab/minute-editor/minute-editor.tsx

@@ -0,0 +1,2 @@
+export const citationRegexWithSpace = /(\s?)\[(\d+)(?:-\d+)?\]/g
+export const citationRegex = /\[(\d+)(?:-\d+)?\]/

frontend/lib/citationRegex.ts

@@ -1,3 +1,4 @@
+import { citationRegexWithSpace } from '@/lib/citationRegex'
 import { DialogueEntry } from '@/lib/client'
 import { saveAs } from 'file-saver'
 import { asBlob } from 'html-docx-js-typescript'
@@ -110,19 +111,12 @@ async function convertHTMLToWordAndDownload(
   saveAs(blob, fileName)
 }
 
-export function removeCitations(text: string): string {
-  const citationRegex = /\s*\[\d+\](\s*\.)?/g
-  return text.replace(citationRegex, (match, punctuation) => {
-    return punctuation ? '.' : ''
-  })
-}
-
 async function convertAIMinutesToWordDoc(
   html: string,
   transcript: DialogueEntry[],
   fileName: string = 'document.docx'
 ): Promise<void> {
-  const cleanedHTML = removeCitations(html)
+  const cleanedHTML = html.replace(citationRegexWithSpace, '')
   await convertHTMLToWordAndDownload(cleanedHTML, transcript, fileName)
 }
 

frontend/lib/download-word-doc.ts

@@ -1,17 +1,18 @@
+import { citationRegex } from '@/lib/citationRegex'
+
 export function linkCitations(content: string): string {
   const parts: string[] = []
-  const regex = /\[(\d+)\]/g
   let lastIndex = 0
   let match: RegExpExecArray | null
-
+  const regex = new RegExp(citationRegex, 'g')
   while ((match = regex.exec(content)) !== null) {
     const idx = match.index
     if (idx > lastIndex) {
       parts.push(content.slice(lastIndex, idx))
     }
     const citationNumber = parseInt(match[1], 10)
     const text = match[0]
-    parts.push(`[[${citationNumber}]](${citationNumber})`)
+    parts.push(`[${text}](${citationNumber})`)
     lastIndex = idx + text.length
   }
   if (lastIndex < content.length) {
@@ -32,7 +33,6 @@ export function CitationContent({
   const regex = /^\d+$/
   let match: RegExpExecArray | null
   match = regex.exec(href || '')
-
   if (match !== null) {
     const idx = match.index
     const text = match[0]
@@ -46,7 +46,7 @@ export function CitationContent({
           onCitationClick(citationNumber, rect)
         }}
       >
-        [{text}]
+        {linkChildren}
       </span>
     )
   }