Merge pull request #6 from HealthRex/feature/populate-json-templates

Feature/populate json templates

Author: ynodir

package-lock.json

@@ -3,24 +3,20 @@ You must strictly follow the Stanford template structure wherever applicable.
 Given:
 1. A clinical question from a PCP
 2. Patient case details
-3. Access to Stanford templates
+3. Stanford JSON template
 Generate JSON object with the following fields:
 - Field: specialistSummary (3-4 lines):
   - Summary of the relevant case details for a specialist who is reviewing the case
   - Focus on patient demographics and key clinical context, relevant history/findings, especially findings that impact management
   - Summarize existing data only; do not include recommendations or suggestions
-- Field: templateSelectionProcess:
-  - First, examine ONLY the templates in the following Stanford document: {{TemplateGoogleDocLink}}
-  - Mention ALL available templates in the above document and use step-by-step logic to justify the most appropriate selection.
 - Field: basicPatientSummary (return as array):
   - Patient age
   - Patient gender
   - Comma separated string of past medical history
   - Comma separated string of key medications
-- Field: populatedTemplate (STRICT FORMAT AND TEMPLATE ADHERENCE AND LAB VALUE INCLUSION REQUIRED):
-  - Respond with JSON array and put each filled template field into its own JSON object, e.g., {"field": "filled template field name", "value": "filled template field value"} as an item of enclosing array.
-  - Maintain original order of sections and questions without modifications.
-  - Populate each field with case-specific information.
+- Field: populatedTemplate:
+  - Populate "***" values of Stanford JSON template with case-specified information
+  - Strictly follow input Stanford JSON format
   - For lab results, explicitly list numeric values instead of checkboxes. Example:
     * Thyroid Stimulating Hormone (TSH): 0.02 mIU/L
     * Free Thyroxine (Free T4): 2.1 ng/dL

resources/prompt.txt

@@ -57,20 +57,26 @@ export class AppService {
   async postReferralQuestion(
     request: ReferralRequest,
   ): Promise<ReferralResponse> {
-    const systemPrompt = await this.selectSystemPrompt(request);
-    const llmResponse = await this.queryLLM<LLMResponse>(
+    const bestTemplate: string =
+      await this.templateSelectorService.selectBestTemplate(request.question);
+    const systemPrompt: string = this.selectSystemPrompt(bestTemplate);
+
+    const llmResponse: LLMResponse = await this.queryLLM<LLMResponse>(
       systemPrompt,
       [
         'Clinical question: ' + request.question,
         'Patient notes: ' + request.clinicalNotes,
       ],
-      llmResponseSchema,
+      llmResponseSchema(JSON.parse(bestTemplate)),
+    );
+    const pathwayResponse: SpecialistAIResponse = await this.queryPathway(
+      request,
+      llmResponse,
     );
-    const pathwayResponse = await this.queryPathway(request, llmResponse);
     const response: ReferralResponse = llmResponse as ReferralResponse;
     response.specialistAIResponse = pathwayResponse;
 
-    this.logger.debug(response);
+    this.logger.debug(JSON.stringify(response, null, 2));
 
     return response;
   }
@@ -79,15 +85,17 @@ export class AppService {
     request: ReferralRequest,
     session: Record<string, any>,
   ): Promise<Observable<{ data: ReferralResponse }>> {
-    const systemPrompt = await this.selectSystemPrompt(request);
+    const bestTemplate: string =
+      await this.templateSelectorService.selectBestTemplate(request.question);
+    const systemPrompt: string = this.selectSystemPrompt(bestTemplate);
     const llmResponseObservable: Observable<ReferralResponse> =
       this.queryLLMStreamed<ReferralResponse>(
         systemPrompt,
         [
           'Clinical question: ' + request.question,
           'Patient notes: ' + request.clinicalNotes,
         ],
-        llmResponseSchema,
+        llmResponseSchema(JSON.parse(bestTemplate)),
       );
 
     return new Observable((subscriber) => {
@@ -100,7 +108,7 @@ export class AppService {
           // reset Pathway conversation history on new referral request
           session[SessionKeys.PREVIOUS_PATHWAY_CONVERSATIONS] = [];
 
-          this.logger.debug('LLM partial response: ', next);
+          this.logger.debug('LLM partial response: ', JSON.stringify(next));
           subscriber.next({ data: next });
         },
         error: (reason) => {
@@ -227,6 +235,19 @@ export class AppService {
     return this.queryLLM<string[]>(systemPrompt, request, responseSchema);
   }
 
+  private selectSystemPrompt(bestTemplate: string) {
+    if (bestTemplate.length > 0) {
+      return fs
+        .readFileSync(join(process.cwd(), systemPromptFilePath))
+        .toString();
+    } else {
+      // no template selected
+      return fs
+        .readFileSync(join(process.cwd(), systemPromptWithoutTemplatesFilePath))
+        .toString();
+    }
+  }
+
   private async queryLLM<T>(
     systemPrompt: string,
     messages: string[],
@@ -310,25 +331,6 @@ export class AppService {
     );
   }
 
-  private async selectSystemPrompt(request: ReferralRequest) {
-    const bestTemplate = await this.templateSelectorService.selectBestTemplate(
-      request.question,
-    );
-    this.logger.debug('bestTemplate', bestTemplate);
-
-    if (bestTemplate) {
-      return fs
-        .readFileSync(join(process.cwd(), systemPromptFilePath))
-        .toString()
-        .replace('{{TemplateGoogleDocLink}}', bestTemplate);
-    } else {
-      // no template selected
-      return fs
-        .readFileSync(join(process.cwd(), systemPromptWithoutTemplatesFilePath))
-        .toString();
-    }
-  }
-
   private selectModel(): LanguageModelV1 {
     switch (String(process.env.AI_PROVIDER).toUpperCase() as AIProvider) {
       case AIProvider.Claude:

src/app.service.ts

@@ -1,5 +1,6 @@
 import { ApiProperty } from '@nestjs/swagger';
 import { z } from 'zod';
+import { jsonToZod } from '../validation/jsonToZodSchema';
 
 export class LLMResponse {
   @ApiProperty({ description: 'Specialist summary response' })
@@ -24,13 +25,14 @@ export class LLMResponse {
   populatedTemplate: Record<string, string>[];
 }
 
-export const llmResponseSchema = z.object({
-  specialistSummary: z.string(),
-  templateSelectionProcess: z.string(),
-  basicPatientSummary: z.array(
-    z.object({ field: z.string(), value: z.string() }),
-  ),
-  populatedTemplate: z.array(
-    z.object({ field: z.string(), value: z.string() }),
-  ),
-});
+export const llmResponseSchema = (
+  templateJson: any,
+): z.ZodObject<any, any, any, any, any> => {
+  return z.object({
+    specialistSummary: z.string(),
+    basicPatientSummary: z.array(
+      z.object({ field: z.string(), value: z.string() }),
+    ),
+    populatedTemplate: jsonToZod(templateJson),
+  });
+};

src/models/llmResponse.ts

@@ -4,6 +4,8 @@ export class SelectBestTemplateResponse {
   question: string;
   @Expose({ name: 'suggested_template' })
   suggestedTemplate: string;
+  @Expose({ name: 'suggested_template_json' })
+  suggestedTemplateJson: string;
   @Expose({ name: 'similarity_scores' })
   similarityScores: object;
 }

src/template-selector/models/selectBestTemplateResponse.ts

@@ -1,77 +1,24 @@
 import { Injectable, Logger } from '@nestjs/common';
 import { HttpService } from '@nestjs/axios';
-import { AxiosError, AxiosResponse } from 'axios';
-import { catchError, lastValueFrom, map, of } from 'rxjs';
+import { AxiosError } from 'axios';
+import { catchError, lastValueFrom, map } from 'rxjs';
 import { SelectBestTemplateRequest } from './models/selectBestTemplateRequest';
 import { SelectBestTemplateResponse } from './models/selectBestTemplateResponse';
 import { plainToInstance } from 'class-transformer';
 
-// TODO would passing drive.google.com/file/.../view link work as well?
-const templateGoogleDocLinks: { [key: string]: string } = {
-  'Allergy eConsult Checklists _extracted.txt':
-    'https://docs.google.com/document/d/1JtFiX31BDZWt0zt0TMKvzLEUxZO5OYlO/edit',
-  'Cardiology eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1oHN8mjupLrk7QVoH-s2hNNM3V4qCM-Cx/view',
-  'Chemical Dependency eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/14Alx6GLZWsl5_D8c57geJSCdWitunPpr/view',
-  'Dermatology eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1Bc695IqIwT65sSz4uxUJraKfN3BxGVOG/view',
-  'Endocrinology eConsult Checklists FINAL 4.19.22_extracted.txt':
-    'https://docs.google.com/document/d/1tA1bkGYlnoFq2OXKMfP996ZFGoSH0DRO/edit',
-  'ENT eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1Cof4fRSaH5HbMmr_O8mK3xrb-2YElpov/view',
-  'Gastroenterology eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1KzIx98EJlN8AGlHaYDBHQ0T4BEV9BvX4/view',
-  'Gynecology eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1iHzbpO-Kmm98z8XxMXtBd191tzluBv0e/view',
-  'Hematology eConsult Checklists FINAL 9.7.22_extracted.txt':
-    'https://drive.google.com/file/d/1nRdw1uIWAAw047snDVXfWGr6AwiwEEeW/view',
-  'Hepatology Oncology eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/14ZSQwWhmqcLETp1nNNst2caPgtghQ4OP/view',
-  'Infectious Disease eConsult Checklists_3.13.24_extracted.txt':
-    'https://drive.google.com/file/d/1W2eXxt3i2U6pdEOiydE4pJLfkZQmQ4qr/view',
-  'Interventional Pulmonology eConsult Checklists _extracted.txt':
-    'https://drive.google.com/file/d/1VCmKnX4wdh9rxE9DOQDxUhVRocZx9gjH/view',
-  'LGBTQ+ eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1UdZrdqu4g7ZDelVJy80pcfZb9JwjQ8yq/view',
-  'Nephrology eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1MhBZ4JgUvHh1vMDsO1b4WXr14zFaaMTk/view',
-  'Neurology eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1yRjPXD90c3aU5Oz7bjt1keQecIo0yEx5/view',
-  'Oncology trial eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1o8VVVuUgSZhLjwCzc-2u3Va0MWOAgFUD/view',
-  'Orthopedics eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1rw0CvTxyExXlSVIhYW_BNCcLtrGu5Yh1/view',
-  'Pain Medicine eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1Sj-riOwZ9ow6JsUh1zvloyuAFl5pJV20/view',
-  'Psychiatry eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1cnJ3MZFkt5UdSpXdLAGBRB6FnNz5tqGN/view',
-  'Pulmonology eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/11qKl1vkOIsFGxYTAuMU6Yo57mrjQqH56/view',
-  'Rheumatology eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1sNB3Ls_RcYoBKjuV3YspAg9h1H30oSCB/view',
-  'SHC Referral Templates Checklists compiled DRAFT_extracted.txt':
-    'https://drive.google.com/file/d/1uQ7GfwsuHb5GVmb7aU0B_VLC_tTxJe7T/view',
-  'Sleep Medicine eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1qdJh1mSNo8OPRsGtOpdkxh8vpeiPS1Pk/view',
-  'Urology eConsult Checklists_extracted.txt':
-    'https://drive.google.com/file/d/1dX80rZgPpzZGMbhdjpqBsHhCJkM3NInS/view',
-};
-
 @Injectable()
 export class TemplateSelectorService {
-  private readonly logger = new Logger(TemplateSelectorService.name);
+  private readonly logger: Logger = new Logger(TemplateSelectorService.name);
 
   constructor(private readonly httpService: HttpService) {}
 
   async selectBestTemplate(clinicalQuestion: string): Promise<string> {
     const request: SelectBestTemplateRequest = new SelectBestTemplateRequest();
     request.question = clinicalQuestion;
 
-    const { data } = await lastValueFrom(
+    const suggestedTemplateJson: string = await lastValueFrom(
       this.httpService
         .post<SelectBestTemplateResponse, SelectBestTemplateRequest>(
-          // TODO get-template is wrong name for POST method - should be smth like select-best-template
           String(process.env.EMBEDDINGS_API) + '/select-best-template',
           request,
         )
@@ -82,19 +29,14 @@ export class TemplateSelectorService {
               SelectBestTemplateResponse,
               response.data,
             );
-            return response;
+            return response.data.suggestedTemplateJson;
           }),
         )
         .pipe(
           catchError((error: AxiosError) => {
             if (error.status === 404) {
               // no template matched
-              return of(
-                error.response as AxiosResponse<
-                  SelectBestTemplateResponse,
-                  SelectBestTemplateRequest
-                >,
-              );
+              return '';
             }
 
             // TODO handle 404 error - need to send different prompt to LLM asking it to generate template by itself
@@ -107,8 +49,8 @@ export class TemplateSelectorService {
           }),
         ),
     );
-    this.logger.debug('data', data);
+    this.logger.debug('suggestedTemplateJson', suggestedTemplateJson);
 
-    return templateGoogleDocLinks[data.suggestedTemplate];
+    return suggestedTemplateJson;
   }
 }

src/template-selector/template-selector.service.ts

@@ -0,0 +1,49 @@
+import { z } from 'zod';
+import { ZodTypeAny } from 'zod/lib/types';
+
+// inspired from https://github.com/colinhacks/zod/discussions/585
+export const jsonToZod = (obj: any): z.ZodTypeAny => {
+  const parse = (obj: any): z.ZodTypeAny => {
+    switch (typeof obj) {
+      case 'string':
+        return z.string();
+      case 'number':
+        return z.number();
+      case 'bigint':
+        return z.number().int();
+      case 'boolean':
+        return z.boolean();
+      case 'object': {
+        if (Array.isArray(obj)) {
+          const options = obj
+            .map(parse)
+            .reduce(
+              (acc: z.ZodTypeAny[], curr) =>
+                acc.includes(curr) ? acc : [...acc, curr],
+              [],
+            );
+          if (options.length === 1) {
+            return z.array(options[0]);
+          } else if (options.length > 1) {
+            return z.array(z.union([options[0], options[1], ...options]));
+          } else {
+            return z.array(z.unknown());
+          }
+        }
+        const zodObj: Map<string, ZodTypeAny> = new Map(
+          Object.entries(obj as object).map(([k, v]) => [k, parse(v)]),
+        );
+        return z.object(Object.fromEntries(zodObj));
+      }
+      case 'undefined':
+        return z.undefined();
+      case 'function':
+        return z.function();
+      case 'symbol':
+      default:
+        return z.unknown();
+    }
+  };
+
+  return parse(obj);
+};