From 56f6637e9d9052c4bc7724d2713573ac67141631 Mon Sep 17 00:00:00 2001
From: "A.J. Shulman" <Shulman.aj@gmail.com>
Date: Mon, 22 Jul 2024 13:20:36 -0400
Subject: works decently well

---
 src/client/views/nodes/ChatBox/Agent.ts            |  4 +-
 src/client/views/nodes/ChatBox/AnswerParser.ts     | 70 ++++++++--------------
 .../views/nodes/ChatBox/MessageComponent.tsx       |  2 +-
 src/client/views/nodes/ChatBox/prompts.ts          |  4 +-
 src/client/views/nodes/ChatBox/tools/RAGTool.ts    | 52 ++++++----------
 5 files changed, 47 insertions(+), 85 deletions(-)

(limited to 'src')

diff --git a/src/client/views/nodes/ChatBox/Agent.ts b/src/client/views/nodes/ChatBox/Agent.ts
index ca1b5c60c..b9d137270 100644
--- a/src/client/views/nodes/ChatBox/Agent.ts
+++ b/src/client/views/nodes/ChatBox/Agent.ts
@@ -22,7 +22,7 @@ export class Agent {
     private _summaries: () => string;
 
     constructor(_vectorstore: Vectorstore, summaries: () => string, history: () => string) {
-        this.client = new OpenAI({ apiKey: process.env.OPENAI_KEY, dangerouslyAllowBrowser: true });
+        this.client = new OpenAI({ apiKey: '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', dangerouslyAllowBrowser: true });
         this.vectorstore = _vectorstore;
         this._history = history;
         this._summaries = summaries;
@@ -34,7 +34,7 @@ export class Agent {
         };
     }
 
-    async askAgent(question: string, maxTurns: number = 10): Promise<string> {
+    async askAgent(question: string, maxTurns: number = 20): Promise<string> {
         console.log(`Starting query: ${question}`);
         this.messages.push({ role: 'user', content: question });
         const chatHistory = this._history();
diff --git a/src/client/views/nodes/ChatBox/AnswerParser.ts b/src/client/views/nodes/ChatBox/AnswerParser.ts
index 4b6c817fd..dd7ec3499 100644
--- a/src/client/views/nodes/ChatBox/AnswerParser.ts
+++ b/src/client/views/nodes/ChatBox/AnswerParser.ts
@@ -17,71 +17,49 @@ export class AnswerParser {
         }
 
         const rawTextContent = answerMatch[1].trim();
-        let textContent: AssistantMessage['content'] = [];
+        let content: AssistantMessage['content'] = [];
         let citations: Citation[] = [];
         let contentIndex = 0;
 
         // Parse citations
         let citationMatch;
+        const citationMap = new Map<string, string>();
         while ((citationMatch = citationRegex.exec(rawTextContent)) !== null) {
             const [_, index, chunk_id, type, direct_text] = citationMatch;
+            const citation_id = uuid();
+            citationMap.set(index, citation_id);
             citations.push({
                 direct_text: direct_text.trim(),
                 type: getChunkType(type),
                 chunk_id,
-                citation_id: uuid(),
+                citation_id,
             });
         }
 
-        // Parse text content (normal and grounded)
-        let lastIndex = 0;
-        let matches = [];
+        // Parse grounded text content
+        const parseGroundedText = (text: string): AssistantMessage['content'] => {
+            const result: AssistantMessage['content'] = [];
+            let lastIndex = 0;
+            let match;
 
-        // Find all grounded text matches
-        let groundedTextMatch;
-        while ((groundedTextMatch = groundedTextRegex.exec(rawTextContent)) !== null) {
-            matches.push({
-                type: 'grounded',
-                index: groundedTextMatch.index,
-                length: groundedTextMatch[0].length,
-                citationIndexes: groundedTextMatch[1],
-                text: groundedTextMatch[2],
-            });
-        }
-
-        // Sort matches by their index in the original text
-        matches.sort((a, b) => a.index - b.index);
-
-        // Process normal and grounded text in order
-        for (let i = 0; i <= matches.length; i++) {
-            const currentMatch = matches[i];
-            const nextMatchIndex = currentMatch ? currentMatch.index : rawTextContent.length;
-
-            // Add normal text before the current grounded text (or end of content)
-            if (nextMatchIndex > lastIndex) {
-                const normalText = rawTextContent.slice(lastIndex, nextMatchIndex).trim();
-                if (normalText) {
-                    textContent.push({
-                        index: contentIndex++,
-                        type: TEXT_TYPE.NORMAL,
-                        text: normalText,
-                        citation_ids: null,
-                    });
-                }
-            }
+            while ((match = groundedTextRegex.exec(text)) !== null) {
+                const [fullMatch, citationIndex, groundedText] = match;
+                const citation_ids = citationIndex.split(',').map(index => citationMap.get(index) || '');
 
-            // Add grounded text if there's a match
-            if (currentMatch) {
-                const citationIds = currentMatch.citationIndexes.split(',').map(index => citations[parseInt(index) - 1].citation_id);
-                textContent.push({
+                result.push({
                     index: contentIndex++,
                     type: TEXT_TYPE.GROUNDED,
-                    text: currentMatch.text.trim(),
-                    citation_ids: citationIds,
+                    text: groundedText.trim(),
+                    citation_ids,
                 });
-                lastIndex = currentMatch.index + currentMatch.length;
+
+                lastIndex = match.index + fullMatch.length;
             }
-        }
+
+            return result;
+        };
+
+        content = parseGroundedText(rawTextContent);
 
         let followUpQuestions: string[] = [];
         if (followUpQuestionsMatch) {
@@ -94,7 +72,7 @@ export class AnswerParser {
 
         const assistantResponse: AssistantMessage = {
             role: ASSISTANT_ROLE.ASSISTANT,
-            content: textContent,
+            content,
             follow_up_questions: followUpQuestions,
             citations,
         };
diff --git a/src/client/views/nodes/ChatBox/MessageComponent.tsx b/src/client/views/nodes/ChatBox/MessageComponent.tsx
index fd7c445c5..d24a55d23 100644
--- a/src/client/views/nodes/ChatBox/MessageComponent.tsx
+++ b/src/client/views/nodes/ChatBox/MessageComponent.tsx
@@ -24,7 +24,7 @@ const MessageComponentBox: React.FC<MessageComponentProps> = function ({ message
                         if (!citation) return null;
                         return (
                             <button
-                                key={idx}
+                                key={i + idx}
                                 className="citation-button"
                                 onClick={() => onCitationClick(citation)}
                                 style={{
diff --git a/src/client/views/nodes/ChatBox/prompts.ts b/src/client/views/nodes/ChatBox/prompts.ts
index 4109ffc17..b2883f71a 100644
--- a/src/client/views/nodes/ChatBox/prompts.ts
+++ b/src/client/views/nodes/ChatBox/prompts.ts
@@ -13,7 +13,7 @@ export function getReactPrompt(tools: Tool[], summaries: () => string, chatHisto
     You operate in a loop of Thought, Action, (STOP), *Action Rules*, Action Input, (STOP), *Observation*, and Answer. Each iteration is numbered (step1, step2, etc.). 
 
     2. Response Structure:
-    a. Enclose each step in numbered XML tags: <step1>, <step2>, etc.
+    a. Enclose each step in numbered XML tags: <step1>, <step2>, etc. (the number of steps will vary for each situation)
     b. Within each step, use the following tags as needed:
         <thought> - Your reasoning process
         <action> - The tool you choose to use
@@ -31,7 +31,7 @@ export function getReactPrompt(tools: Tool[], summaries: () => string, chatHisto
         c.4) Based on the action rules, provide tool parameters in <action_input>.
         c.5) End your response again and wait for the observation from the tool.
     d. If no tool is needed, use the 'no_tool' action but still follow the same structure.
-    e. Use the retrieval (RAG) tool ANYTIME the question may potentially (even if you are not sure) relate to one of the user's documents. Here are the summaries of the user's documents:
+    e. !!!Use the retrieval (RAG) tool ANYTIME the question may potentially (even if you are not sure) relate to one of the user's documents. Here are the summaries of the user's documents:
             ${summaries()}
     f. Based on observations or your knowledge, formulate your answer.
     g. Provide the final answer in the <answer> tag, including follow-up questions.
diff --git a/src/client/views/nodes/ChatBox/tools/RAGTool.ts b/src/client/views/nodes/ChatBox/tools/RAGTool.ts
index 4b29d6bce..c7175326c 100644
--- a/src/client/views/nodes/ChatBox/tools/RAGTool.ts
+++ b/src/client/views/nodes/ChatBox/tools/RAGTool.ts
@@ -22,7 +22,7 @@ export class RAGTool extends BaseTool<{ hypothetical_document_chunk: string }> {
 
             1. Overall Structure:
             <answer>
-                [Main content with nested grounded_text tags]
+                [Main content with grounded_text tags interspersed with normal plain text (information that is not derived from chunks' information)]
                 <citations>
                 [Individual citation tags]
                 </citations>
@@ -33,21 +33,8 @@ export class RAGTool extends BaseTool<{ hypothetical_document_chunk: string }> {
 
             2. Grounded Text Tag Structure:
             - Basic format: 
-                <grounded_text citation_index="[index number(s)]">
-                [Your generated text based on chunk information]
-                </grounded_text>
-
-            - Nested format:
-                <grounded_text citation_index="[index number(s)]">
-                [General information]
-                <grounded_text citation_index="[index number(s)]">
-                    [More specific information]
-                </grounded_text>
-                </grounded_text>
-
-            - Multiple citation indices:
-                <grounded_text citation_index="1,2,3">
-                [Information synthesized from multiple chunks]
+                <grounded_text citation_index="[citation index number(s)]">
+                [Your generated text based on information from a subset of a chunk (a citation's direct text)]
                 </grounded_text>
 
             3. Citation Tag Structure:
@@ -58,12 +45,11 @@ export class RAGTool extends BaseTool<{ hypothetical_document_chunk: string }> {
 
             4. Detailed Grounded Text Guidelines:
                 a. Wrap all information derived from chunks in grounded_text tags.
-                b. Nest grounded_text tags when presenting hierarchical or increasingly specific information or when a larger section of generated text is best grounded by one subset of a chunk and smaller sections of that generated text are best grounded by other subsets of either the same or different chunk(s).
-                c. Use a single grounded_text tag for closely related information that references the same citation (subset of text from a chunk).
-                d. Combine multiple citation indices for synthesized information from multiple citations.
-                e. Ensure every grounded_text tag has at least one corresponding citation.
-                f. Grounded text can be as short as a few words or as long as several sentences.
-                d. Avoid overlapping grounded_text tags; instead, use nesting or sequential tags.
+                b. DO NOT PUT ANYTHING THAT IS NOT DIRECTLY DERIVED FROM INFORMATION FROM CHUNKS (EITHER IMAGE, TABLE, OR TEXT) IN GROUNDED_TEXT TAGS.
+                c. Use a single grounded_text tag for suquential and closely related information that references the same citation. If other citations' information are used sequentially, create new grounded_text tags.
+                d. Ensure every grounded_text tag has up to a few corresponding citations (should not be more than 3 and only 1 is fine). Multiple citation indices should be separated by commas.
+                e. Grounded text can be as short as a few words or as long as several sentences.
+                f. Avoid overlapping or nesting grounded_text tags; instead, use sequential tags.
 
             5. Detailed Citation Guidelines:
                 a. Create a unique citation for each distinct piece of information from the chunks that is used to support grounded_text.
@@ -71,35 +57,33 @@ export class RAGTool extends BaseTool<{ hypothetical_document_chunk: string }> {
                 c. Specify the correct type: "text", "image", or "table".
                 d. For text chunks, include only the relevant subset of the original text that the grounded_text is based on.
                 e. For image/table chunks, leave the citation content empty.
-                f. One citation can be used for multiple grounded_text tags if they are based on the same information.
-                g. One text chunk can have multiple citations if different parts of the text have different important information.
-                h. !!!DO NOT OVERCITE - only include citations for information that is directly relevant to the grounded_text.
+                f. One citation can be used for multiple grounded_text tags if they are based on the same chunk information.
+                g. !!!DO NOT OVERCITE - only include citations for information that is directly relevant to the grounded_text.
 
             6. Structural Integrity Checks:
                 a. Ensure all opening tags have corresponding closing tags.
                 b. Verify that all grounded_text tags have valid citation_index attributes.
                 c. Check that all cited indices in grounded_text tags have corresponding citations.
-                d. Confirm proper nesting - tags opened last should be closed first.
 
             Example of grounded_text usage:
 
             <answer>
                 <grounded_text citation_index="1,2">
                 Artificial Intelligence (AI) is revolutionizing various sectors, with healthcare experiencing significant transformations in areas such as diagnosis and treatment planning.
-                    <grounded_text citation_index="2,3,4">
-                        In the field of medical diagnosis, AI has shown remarkable capabilities, particularly in radiology. For instance, AI systems have drastically improved mammogram analysis, achieving 99% accuracy at a rate 30 times faster than human radiologists.
-                            <grounded_text citation_index="4">
-                            This advancement not only enhances the efficiency of healthcare systems but also significantly reduces the occurrence of false positives, leading to fewer unnecessary biopsies and reduced patient stress.
-                            </grounded_text>
-                    </grounded_text>
                 </grounded_text>
-
+                <grounded_text citation_index="2,3,4">
+                In the field of medical diagnosis, AI has shown remarkable capabilities, particularly in radiology. For instance, AI systems have drastically improved mammogram analysis, achieving 99% accuracy at a rate 30 times faster than human radiologists.
+                </grounded_text>
+                <grounded_text citation_index="4">
+                This advancement not only enhances the efficiency of healthcare systems but also significantly reduces the occurrence of false positives, leading to fewer unnecessary biopsies and reduced patient stress.
+                </grounded_text>
+            
                 <grounded_text citation_index="5,6">
                 Beyond diagnosis, AI is playing a crucial role in drug discovery and development. By analyzing vast amounts of genetic and molecular data, AI algorithms can identify potential drug candidates much faster than traditional methods.
+                </grounded_text>
                 <grounded_text citation_index="6">
                     This could potentially reduce the time and cost of bringing new medications to market, especially for rare diseases that have historically received less attention due to limited market potential.
                 </grounded_text>
-                </grounded_text>
 
                 [... rest of the content ...]
 
-- 
cgit v1.2.3-70-g09d2