Improved AI Actions: Default OpenAI Key & Tracing (#564)

* update core packadge

* refactor update tags action

* extract ai action to its own function

* refactor findAppoitnmentByPrompt

* add type

* refactor final actions

* chor(dependencies): ugrade langchain and langsmith

* feat(shelly): langsmith config (metadata and runname) added as an example for one ai action

* feat(ai-actions): extracting LLM models to separate lib; apply to categorizeMessage

* chore(shelly): refactor categorizeMessage to serve as an example for other ai actions

* chore(shelly): go back to skipping tests with real OpenAI calls

* chore(shelly): refactor generateMessage

* chore(shelly): go back to skipping real openAI tests

* chore(shelly): refactor summarizeCareFlow

* chore(shelly): refactor summarizeForm ai actions

* chore(shelly): remove unnecessary setups form validatePayload

* feat(shelly): OpenAI key optional, experimental->beta for consistency

* chore(elation): refactor UpdatePatient to match other ai actions

* chor(elation): refactor findFutureAppointment

* chore(elation): refactor findAppointmentsByPrompt

* remove console logs

* feat(elation): remove openAIAPI key from settings

* fix(shelly): fix test

* fix(shelly): skip real world test

* fix(llm): deal with no OpenAI api key in settings

* chore: upgrade extension-core

* feat(lib): add tenant_id for tracing,accept more flexible payload

* chore(ai-actions): remove try-catch error wrappers

* chore(tests): add tenant_id to pathway

* fix(tests): remove OpenAI dependency

* feat(langsmith): add callback to make it easy to hide input/output data when tracing (call will still be registered, only data will not be stored)

* feat(shelly): remove settings (OpenAI Key)

* feat(dependencies): upgrade extensions-core to use org stycth information in Pathway

* feat(langsmith): preserve token usage data when masking to make sure we can still track usage and $

* chor(shelly): go back to skipping real world tests

* fix(tests): add missing params to Pathway

* chor(ai-actions): minor refactoring

---------

Co-authored-by: belmai <[email protected]>

.env.example

@@ -4,4 +4,13 @@ ALGOLIA_ADMIN_KEY=
 # Needed to generate the Healthie SDK
 HEALTHIE_API_URL=
 # Only needed if you want to run real test against the OpenAI API
-OPENAI_TEST_KEY=
+OPENAI_API_KEY=
+
+# Langsmith tracing
+LANGSMITH_TRACING=true
+LANGSMITH_ENDPOINT=https://api.smith.langchain.com
+# For now let's set both to make sure it works for all sdk versions
+LANGSMITH_PROJECT=ai-actions-local
+LANGCHAIN_PROJECT=ai-actions-local
+LANGSMITH_API_KEY=
+

extensions/awell/v1/actions/isPatientEnrolledInCareFlow/isPatientEnrolledInCareFlow.test.ts

@@ -125,6 +125,7 @@ describe('Is patient already enrolled in care flow action', () => {
         pathway: {
           id: 'pathway-instance-id-1',
           definition_id: 'pathway-definition-1',
+          tenant_id: '123',
         },
         fields: {
           pathwayStatus: '', // By default, only active care flows
@@ -167,6 +168,7 @@ describe('Is patient already enrolled in care flow action', () => {
         pathway: {
           id: 'pathway-instance-id-1',
           definition_id: 'pathway-definition-1',
+          tenant_id: '123',
         },
         fields: {
           pathwayStatus: undefined, // By default, only active care flows
@@ -215,6 +217,7 @@ describe('Is patient already enrolled in care flow action', () => {
         pathway: {
           id: 'pathway-instance-id-1',
           definition_id: 'pathway-definition-1',
+          tenant_id: '123',
         },
         fields: {
           pathwayStatus: `${PathwayStatus.Completed}`,
@@ -318,6 +321,7 @@ describe('Is patient already enrolled in care flow action', () => {
         pathway: {
           id: 'pathway-instance-id-1',
           definition_id: 'pathway-definition-1',
+          tenant_id: '123',
         },
         fields: {
           pathwayStatus: `${PathwayStatus.Active}, ${PathwayStatus.Completed}`,

extensions/bland/actions/sendCall/__tests__/sendCall.test.ts

@@ -47,6 +47,7 @@ describe('Bland.ai - Send call', () => {
         pathway: {
           id: 'pathway-id',
           definition_id: 'pathway-definition-id',
+          tenant_id: '123',
         },
         activity: {
           id: 'activity-id',

extensions/bland/actions/sendCallWithPathway/__tests__/sendCallWithPathway.test.ts

@@ -47,6 +47,7 @@ describe('Bland.ai - Send call', () => {
         pathway: {
           id: 'pathway-id',
           definition_id: 'pathway-definition-id',
+          tenant_id: '123',
         },
         activity: {
           id: 'activity-id',

extensions/elation/actions/closeCareGap/closeCareGap.test.ts

@@ -46,6 +46,9 @@ describe('Elation - Close care gap', () => {
     pathway: {
       definition_id: '123',
       id: '123',
+      tenant_id: '123',
+      org_id: '123',
+      org_slug: 'org-slug',
     },
     patient: {
       id: '123',

extensions/elation/actions/createCareGap/createCareGap.test.ts

@@ -50,6 +50,9 @@ describe('Elation - Create care gap', () => {
     pathway: {
       definition_id: '123',
       id: '123',
+      tenant_id: '123',
+      org_id: '123',
+      org_slug: 'org-slug',
     },
     patient: {
       id: '123',

extensions/elation/actions/createReferralOrder/createReferralOrder.test.ts

@@ -53,6 +53,9 @@ describe('Elation - Create referral order', () => {
     pathway: {
       definition_id: '123',
       id: '123',
+      tenant_id: '123',
+      org_id: '123',
+      org_slug: 'org-slug',
     },
     patient: {
       id: '123',

extensions/elation/actions/findAppointmentsByPrompt/findAppointmentsByPrompt.test.ts

@@ -1,58 +1,53 @@
+import { TestHelpers } from '@awell-health/extensions-core'
 import { makeAPIClient } from '../../client'
 import { appointmentsMock } from './__testdata__/GetAppointments.mock'
 import { findAppointmentsByPrompt as action } from './findAppointmentsByPrompt'
-import { TestHelpers } from '@awell-health/extensions-core'
-import { ChatOpenAI } from '@langchain/openai'
 
-jest.mock('../../client', () => ({
-  makeAPIClient: jest.fn().mockImplementation(() => ({
-    findAppointments: jest.fn().mockResolvedValue(appointmentsMock),
-  })),
-}))
+// Mock the client
+jest.mock('../../client')
 
-const mockedSdk = jest.mocked(makeAPIClient)
-
-jest.mock('@langchain/openai', () => {
-  const mockInvoke = jest.fn().mockResolvedValue({
-    appointmentIds: appointmentsMock.map((appointment) => appointment.id),
-    explanation: 'Test explanation',
+// Mock createOpenAIModel
+jest.mock('../../../../src/lib/llm/openai/createOpenAIModel', () => ({
+  createOpenAIModel: jest.fn().mockResolvedValue({
+    model: {
+      pipe: jest.fn().mockReturnValue({
+        invoke: jest.fn().mockResolvedValue({
+          appointmentIds: appointmentsMock.map(a => a.id),
+          explanation: 'Test explanation'
+        })
+      })
+    },
+    metadata: {
+      care_flow_definition_id: 'whatever',
+      care_flow_id: 'test-flow-id',
+      activity_id: 'test-activity-id',
+      tenant_id: '123',
+      org_id: '123',
+      org_slug: 'org-slug'
+    }
   })
+}))
 
-  const mockChain = {
-    invoke: mockInvoke,
-  }
-
-  const mockPipe = jest.fn().mockReturnValue(mockChain)
-
-  const mockChatOpenAI = jest.fn().mockImplementation(() => ({
-    pipe: mockPipe,
-  }))
-
-  return {
-    ChatOpenAI: mockChatOpenAI,
-  }
-})
-
-describe('Elation - Find appointment by type', () => {
-  const {
-    extensionAction: findAppointmentByType,
-    onComplete,
-    onError,
-    helpers,
-    clearMocks,
-  } = TestHelpers.fromAction(action)
+describe('Elation - Find appointments by prompt', () => {
+  const { extensionAction, onComplete, onError, helpers, clearMocks } = 
+    TestHelpers.fromAction(action)
 
   beforeEach(() => {
     clearMocks()
     jest.clearAllMocks()
+
+    const mockAPIClient = makeAPIClient as jest.Mock
+    mockAPIClient.mockImplementation(() => ({
+      findAppointments: jest.fn().mockResolvedValue(appointmentsMock)
+    }))
   })
 
-  test('Should return the correct appointment', async () => {
-    await findAppointmentByType.onEvent({
+  test('Should find the correct appointments', async () => {
+    await extensionAction.onEvent({
       payload: {
         fields: {
-          patientId: 12345, // used to get a list of appointments
-          prompt: 'Find the next appointment for this patient',
+          patientId: 12345,
+          prompt: 'Find all appointments',
         },
         settings: {
           client_id: 'clientId',
@@ -61,24 +56,90 @@ describe('Elation - Find appointment by type', () => {
           password: 'password',
           auth_url: 'authUrl',
           base_url: 'baseUrl',
-          openAiApiKey: 'openaiApiKey',
         },
-      } as any,
+        pathway: {
+          id: 'test-flow-id',
+          definition_id: '123',
+          tenant_id: '123',
+          org_slug: 'test-org-slug',
+          org_id: 'test-org-id'
+        },
+        activity: {
+          id: 'test-activity-id'
+        },
+        patient: {
+          id: 'test-patient-id'
+        }
+      },
       onComplete,
       onError,
       helpers,
     })
 
-    expect(ChatOpenAI).toHaveBeenCalled()
-    expect(mockedSdk).toHaveBeenCalled()
-    expect(onComplete).toHaveBeenCalledWith(
-      expect.objectContaining({
-        data_points: {
-          appointments: JSON.stringify(appointmentsMock),
-          explanation: 'Test explanation',
-          appointmentCountsByStatus: JSON.stringify({ Scheduled: 2 }),
+    expect(onComplete).toHaveBeenCalledWith({
+      data_points: {
+        appointments: JSON.stringify(appointmentsMock),
+        explanation: 'Test explanation',
+        appointmentCountsByStatus: JSON.stringify({ Scheduled: 2 }),
+      },
+      events: [
+        {
+          date: expect.any(String),
+          text: {
+            en: `Found ${appointmentsMock.length} appointments for patient ${12345}`
+          }
+        }
+      ],
+    })
+    expect(onError).not.toHaveBeenCalled()
+  })
+
+  test('Should handle no appointments', async () => {
+    const mockAPIClient = makeAPIClient as jest.Mock
+    mockAPIClient.mockImplementation(() => ({
+      findAppointments: jest.fn().mockResolvedValue([])
+    }))
+
+    await extensionAction.onEvent({
+      payload: {
+        fields: {
+          patientId: 12345,
+          prompt: 'Find all appointments',
+        },
+        settings: {
+          client_id: 'clientId',
+          client_secret: 'clientSecret',
+          username: 'username',
+          password: 'password',
+          auth_url: 'authUrl',
+          base_url: 'baseUrl',
         },
-      }),
-    )
+        pathway: {
+          id: 'test-flow-id',
+          definition_id: '123',
+          tenant_id: '123',
+          org_slug: 'test-org-slug',
+          org_id: 'test-org-id'
+        },
+        activity: {
+          id: 'test-activity-id'
+        },
+        patient: {
+          id: 'test-patient-id'
+        }
+      },
+      onComplete,
+      onError,
+      helpers,
+    })
+
+    expect(onComplete).toHaveBeenCalledWith({
+      data_points: {
+        explanation: 'No appointments found',
+        appointments: JSON.stringify([]),
+        appointmentCountsByStatus: JSON.stringify({}),
+      }
+    })
+    expect(onError).not.toHaveBeenCalled()
   })
 })