(null);
14 |
15 | useEffect(() => {
16 | const fetchAttributes = async () => {
17 | try {
18 | const userAttributesResponse = await fetchUserAttributes();
19 | if (userAttributesResponse) setUserAttributes(userAttributesResponse);
20 | } catch (error) {
21 | console.error('Error fetching user attributes:', error);
22 | setUserAttributes(null);
23 | }
24 | };
25 |
26 | fetchAttributes();
27 | }, []);
28 |
29 | // Pass the object as the value
30 | return (
31 |
32 | {children}
33 |
34 | );
35 | }
36 |
37 | // Custom hook to use the context
38 | export function useUserAttributes() {
39 | const context = useContext(UserContext);
40 | if (context === undefined) {
41 | throw new Error('useUserAttributes must be used within a UserProvider');
42 | }
43 | return context;
44 | }
--------------------------------------------------------------------------------
/src/components/WithAuth.tsx:
--------------------------------------------------------------------------------
1 | import React, { useEffect } from 'react';
2 | import { useAuthenticator } from '@aws-amplify/ui-react';
3 | import { redirect } from 'next/navigation';
4 |
5 | export function withAuth(Component: React.ComponentType
) {
6 | return function AuthProtected(props: P) {
7 | const { authStatus } = useAuthenticator(context => [context.authStatus]);
8 |
9 | useEffect(() => {
10 | if (authStatus === 'unauthenticated') {
11 | redirect('/login')
12 | }
13 | }, [authStatus]);
14 |
15 | if (authStatus === 'authenticated') {
16 | return ;
17 | }
18 |
19 | return null;
20 | };
21 | }
22 |
--------------------------------------------------------------------------------
/src/hero-img.jpg:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/agents4energy/40866cab8ce43e16e87f5b51a084538ca8295044/src/hero-img.jpg
--------------------------------------------------------------------------------
/src/hero-img.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/aws-samples/agents4energy/40866cab8ce43e16e87f5b51a084538ca8295044/src/hero-img.png
--------------------------------------------------------------------------------
/src/logo-small-top-navigation.svg:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/src/styles/chat-ui.module.scss:
--------------------------------------------------------------------------------
1 | @use "@cloudscape-design/design-tokens/index" as awsui;
2 |
3 | .chat_container {
4 | margin-bottom: -40px;
5 | // min-height: calc(100vh - 200px); //96px
6 | // max-height: calc(20vh); //96px
7 | display: flex;
8 | flex-direction: column;
9 | justify-content: space-between;
10 | }
11 |
12 | @media (min-width: 466px) {
13 | .chat_container {
14 | min-height: calc(100vh - 300px);
15 | }
16 | }
17 |
18 | // @media (min-width: 689px) {
19 | // .chat_container {
20 | // min-height: calc(100vh - 60px);
21 | // }
22 | // }
23 |
24 | // @media (min-width: 913px) {
25 | // .chat_container {
26 | // min-height: calc(100vh - 68px);
27 | // }
28 | // }
29 |
30 | .welcome_text {
31 | color: awsui.$color-text-body-default;
32 | font-size: 3rem;
33 | font-weight: bolder;
34 | opacity: 0.4;
35 | }
36 |
37 | .input_container {
38 | position: sticky;
39 | bottom: 0;
40 | padding-bottom: awsui.$space-scaled-l;
41 | padding-top: awsui.$space-scaled-xxl;
42 | background: linear-gradient(
43 | to bottom,
44 | transparent 0%,
45 | awsui.$color-background-container-content 20%
46 | );
47 | margin-left: -8px;
48 | margin-right: -8px;
49 | container-type: inline-size;
50 | container-name: input_container;
51 | }
52 |
53 | .input_textarea_container {
54 | display: grid;
55 | grid-template-columns: 1fr auto;
56 | align-items: center;
57 | gap: 4px;
58 | }
59 |
60 | .input_textarea {
61 | resize: none;
62 | border: none;
63 | padding: 12px 12px 12px 4px;
64 | background-color: transparent;
65 | outline: none;
66 | width: 100%;
67 | height: 100%;
68 | font-size: 1rem;
69 | }
70 |
71 | .btn_chabot_message_copy {
72 | float: right;
73 | }
74 |
75 | .horizontal-split {
76 | flex: 1
77 | }
78 | .markdownTable {
79 | width: 100%;
80 | border-collapse: collapse;
81 | }
82 |
83 | .markdownTable th,
84 | .markdownTable td {
85 | padding: 8px;
86 | text-align: left;
87 | border-bottom: 1px solid #ddd; /* Adds a line below each row */
88 | }
89 |
90 | .markdownTable tr + tr {
91 | border-top: 1px solid #ddd; /* Adds a line above each row */
92 | }
93 |
94 | .hiddenRow {
95 | display: none; /* Hides the row */
96 | }
97 |
98 | .prose {
99 | flex: 1;
100 | overflow-y: auto; /* If you want scrolling when content overflows */
101 | color: unset!important;
102 | }
103 |
104 |
--------------------------------------------------------------------------------
/src/utils/chart-utils.ts:
--------------------------------------------------------------------------------
1 | // import { Chart as ChartJS, Plugin } from 'chart.js';
2 | // import 'chartjs-plugin-datalabels';
3 | // import type { Context } from 'chartjs-plugin-datalabels';
4 |
5 | // interface DataLabel {
6 | // x: number;
7 | // y: number;
8 | // width: number;
9 | // height: number;
10 | // datasetIndex: number;
11 | // index: number;
12 | // }
13 |
14 | // export const CustomDataLabelsPlugin: Plugin = {
15 | // id: 'customDataLabels',
16 | // beforeDraw(chart: ChartJS) {
17 | // const ctx = chart.ctx;
18 | // const dataLabels: DataLabel[] = [];
19 |
20 | // // Get all data labels positions and dimensions
21 | // chart.data.datasets.forEach((dataset, datasetIndex) => {
22 | // const meta = chart.getDatasetMeta(datasetIndex);
23 |
24 | // if (!meta.hidden) {
25 | // meta.data.forEach((element, index) => {
26 | // const { x, y } = element.getCenterPoint();
27 |
28 | // // Get the data label for this point
29 | // const value = dataset.data[index];
30 | // const label = chart.options?.plugins?.datalabels?.formatter?.(
31 | // value,
32 | // { dataIndex: index, dataset, datasetIndex }
33 | // )?.toString() || value?.toString();
34 |
35 | // // Measure text width and height
36 | // const textMetrics = ctx.measureText(label);
37 | // const textWidth = textMetrics.width;
38 | // const fontSize = chart.options?.plugins?.datalabels?.font?.size as number || 12;
39 |
40 | // dataLabels.push({
41 | // x,
42 | // y,
43 | // width: textWidth,
44 | // height: fontSize,
45 | // datasetIndex,
46 | // index
47 | // });
48 | // });
49 | // }
50 | // });
51 |
52 | // // Check for overlaps and adjust positions
53 | // for (let i = 0; i < dataLabels.length; i++) {
54 | // for (let j = i + 1; j < dataLabels.length; j++) {
55 | // if (checkOverlap(dataLabels[i], dataLabels[j])) {
56 | // const meta = chart.getDatasetMeta(dataLabels[j].datasetIndex);
57 | // const datalabelsPlugin = meta.controller.$context.chart.options.plugins.datalabels;
58 |
59 | // if (dataLabels[i].y > dataLabels[j].y) {
60 | // // Move second label up
61 | // datalabelsPlugin.align = 'bottom';
62 | // datalabelsPlugin.offset = 10;
63 | // } else {
64 | // // Move second label down
65 | // datalabelsPlugin.align = 'top';
66 | // datalabelsPlugin.offset = -10;
67 | // }
68 | // }
69 | // }
70 | // }
71 | // }
72 | // };
73 |
74 | // function checkOverlap(label1: DataLabel, label2: DataLabel): boolean {
75 | // return !(
76 | // label1.x + label1.width / 2 < label2.x - label2.width / 2 ||
77 | // label1.x - label1.width / 2 > label2.x + label2.width / 2 ||
78 | // label1.y + label1.height / 2 < label2.y - label2.height / 2 ||
79 | // label1.y - label1.height / 2 > label2.y + label2.height / 2
80 | // );
81 | // }
--------------------------------------------------------------------------------
/src/utils/config.ts:
--------------------------------------------------------------------------------
1 | // import { BedrockAgent } from "@aws-sdk/client-bedrock-agent"
2 | import outputs from '@/../amplify_outputs.json';
3 |
4 | type BaseAgent = {
5 | name: string
6 | samplePrompts: string[]
7 | source: 'bedrockAgent' | 'graphql'
8 | }
9 |
10 | export type BedrockAgent = BaseAgent & {
11 | source: "bedrockAgent"
12 | agentId: string
13 | agentAliasId: string
14 | }
15 |
16 | export type LangGraphAgent = BaseAgent & {
17 | source: "graphql"
18 | invokeFieldName: string
19 | }
20 |
21 | export const defaultAgents: { [key: string]: BaseAgent | BedrockAgent | LangGraphAgent } = {
22 | PlanAndExecuteAgent: {
23 | name: `Production Agent`,
24 | source: `graphql`,
25 | samplePrompts: [
26 | `This morning well with API number 30-045-29202 stopped producing gas with indication of a hole in tubing.
27 | Make a table of all operational events found in the well files.
28 | Query all historic monthly production rates and make a plot with both the event and production data.
29 | Estimate the value of the well's remaining production.
30 | Write a procedure to repair the well, estimate the cost of the repair, and calculate financial metrics.
31 | Make an executive report about repairing the well with detailed cost and procedure data.
32 | Use the ai role for all steps.
33 | `.replace(/^\s+/gm, ''),
34 | `Search the well files for the well with API number 30-045-29202 to make a table with type of operation (drilling, completion, workover, plugging, other), text from the report describing operational details, and document title.
35 | Also execute a sql query to get the total monthly oil, gas and water production from this well.
36 | Create a plot with both the event data and the production data. `.replace(/^\s+/gm, ''), //This trims the white space at the start of each line
37 | `Plot the total monthly oil, gas, and water production since 1900 for the well with API number 30-045-29202`,
38 | `Which form of artifical lift best matches my personality?`
39 | ]
40 | },
41 | MaintenanceAgent: {
42 | name: "Maintenance Agent",
43 | source: "bedrockAgent",
44 | agentId: outputs.custom.maintenanceAgentId,
45 | agentAliasId: outputs.custom.maintenanceAgentAliasId,
46 | samplePrompts: [
47 | "How many tanks are in my biodiesel unit?",
48 | "In September 2024, what are a few key incidents and actions taken at the biodiesel unit?",
49 | ],
50 | } as BedrockAgent,
51 | RegulatoryAgent: {
52 | name: "Regulatory Agent",
53 | source: "bedrockAgent",
54 | agentId: outputs.custom.regulatoryAgentId,
55 | agentAliasId: outputs.custom.regulatoryAgentAliasId,
56 | samplePrompts: [
57 | "What are the requirements for fugitive emissions monitoring and reporting in the U.S.?",
58 | "What are the requirements for decomissioning an offshore oil well in Brazil?",
59 | ],
60 | } as BedrockAgent,
61 | PetrophysicsAgent: {
62 | name: "Petrophysics Agent",
63 | source: "bedrockAgent",
64 | agentId: outputs.custom.petrophysicsAgentId,
65 | agentAliasId: outputs.custom.petrophysicsAgentAliasId,
66 | samplePrompts: [
67 | "Give me a summary fluid substitution modeling",
68 | "Give me the inputs of Gassmann equation",
69 | "What are AVO classes?",
70 | "Calculate the intercept and gradient value of the wet sandstone with vp=3.5 km/s, vs=1.95 km/s, bulk density=2.23 gm/cc when it is overlain by a shale? Determine the AVO class.",
71 | "A wet sandstone has vp=3.5 km/s, vs=1.95 km/s, bulk density=2.23 gm/cc. What are the expected seismic velocities of the sandstone if the desired fluid saturation is 80% oil? Use standard assumptions."
72 | ],
73 | } as BedrockAgent
74 | }
--------------------------------------------------------------------------------
/src/utils/date-utils.ts:
--------------------------------------------------------------------------------
1 | import { format, parseISO } from 'date-fns';
2 |
3 | export const formatDate = (dateString: string): string => {
4 | const date = parseISO(dateString);
5 | return format(date, 'MMM d, yyyy HH:mm:ss');
6 | };
--------------------------------------------------------------------------------
/src/utils/rateLimiter.ts:
--------------------------------------------------------------------------------
1 | import { SSMClient, GetParameterCommand } from '@aws-sdk/client-ssm';
2 |
3 | export class RateLimiter {
4 | private static instance: RateLimiter;
5 | private lastInvocationTime: number | null = null;
6 | private rateLimitSeconds: number;
7 |
8 | private static async initializeRateLimit(): Promise {
9 | const ssm = new SSMClient({});
10 | const parameterPath = 'AGENT_RATE_LIMIT';
11 | console.log(`RateLimiter: Attempting to fetch rate limit from Parameter Store at path: ${parameterPath}`);
12 | const startTime = Date.now();
13 |
14 | try {
15 | const command = new GetParameterCommand({
16 | Name: parameterPath,
17 | WithDecryption: false
18 | });
19 | const result = await ssm.send(command);
20 | const fetchDuration = Date.now() - startTime;
21 |
22 | if (!result.Parameter?.Value) {
23 | console.log('RateLimiter: No rate limit value found in Parameter Store');
24 | console.log(`RateLimiter: Parameter Store fetch took ${fetchDuration}ms`);
25 | return 0;
26 | }
27 |
28 | const rateLimit = parseInt(result.Parameter.Value);
29 | if (isNaN(rateLimit)) {
30 | console.error(`RateLimiter: Invalid rate limit value: ${result.Parameter.Value}. Using default rate limit of 0 seconds.`);
31 | return 0;
32 | }
33 |
34 | console.log(`RateLimiter: Found rate limit value: ${rateLimit} seconds`);
35 | console.log(`RateLimiter: Parameter Store fetch took ${fetchDuration}ms`);
36 | return rateLimit;
37 | } catch (error) {
38 | const fetchDuration = Date.now() - startTime;
39 | console.error('RateLimiter: Error fetching rate limit from Parameter Store:', error);
40 | console.log(`RateLimiter: Parameter Store fetch took ${fetchDuration}ms`);
41 | return 0;
42 | }
43 | }
44 |
45 | private constructor() {
46 | this.rateLimitSeconds = 0;
47 | }
48 |
49 | public static async getInstance(): Promise {
50 | if (!RateLimiter.instance) {
51 | RateLimiter.instance = new RateLimiter();
52 | RateLimiter.instance.rateLimitSeconds = await RateLimiter.initializeRateLimit();
53 | }
54 | return RateLimiter.instance;
55 | }
56 |
57 | public async waitForRateLimit(): Promise {
58 | if (this.rateLimitSeconds <= 0) {
59 | console.log('RateLimiter: Skipping rate limiting as it is not configured.');
60 | return;
61 | }
62 |
63 | const now = Date.now();
64 | if (this.lastInvocationTime !== null) {
65 | const timeSinceLastInvocation = now - this.lastInvocationTime;
66 | const waitTime = Math.max(0, (this.rateLimitSeconds * 1000) - timeSinceLastInvocation);
67 |
68 | if (waitTime > 0) {
69 | console.log(`RateLimiter: Pausing for ${waitTime}ms due to rate limit threshold being met`);
70 | await new Promise(resolve => setTimeout(resolve, waitTime));
71 | } else {
72 | console.log(`RateLimiter: Below threshold (last call ${timeSinceLastInvocation}ms ago)`);
73 | }
74 | }
75 |
76 | this.lastInvocationTime = now;
77 | }
78 | }
--------------------------------------------------------------------------------
/src/utils/types.ts:
--------------------------------------------------------------------------------
1 | import type { Schema } from '../../amplify/data/resource';
2 | export type Message = Schema["ChatMessage"]["createType"] & {
3 | previousTrendTableMessage?: Schema["ChatMessage"]["createType"];
4 | previousEventTableMessage?: Schema["ChatMessage"]["createType"];
5 | };
6 |
7 |
8 |
9 | export type messageContentType = 'ai' | 'tool_markdown' | 'tool_json' | 'tool_table_trend' | 'tool_table_events' | 'tool_plot'
10 |
11 | export type ToolMessageContentType = {
12 | messageContentType: messageContentType;
13 | };
14 |
15 | // export type Message = {
16 | // content: string;
17 | // owner?: string;
18 | // role: string;
19 | // createdAt?: string;
20 | // trace?: string
21 | // tool_name?: string;
22 | // tool_call_id?: string;
23 | // tool_calls?: string;
24 | // chatSessionId?: string;
25 | // }
--------------------------------------------------------------------------------
/src/utils/ui-utils.tsx:
--------------------------------------------------------------------------------
1 | import { ComponentProps } from 'react';
2 | import { Message } from '@/utils/types'
3 |
4 | export const jsonParseHandleError = (jsonString: string) => {
5 | try {
6 | return JSON.parse(jsonString)
7 | } catch {
8 | console.warn(`Could not parse string: ${jsonString}`)
9 | }
10 | }
11 |
12 | export const combineAndSortMessages = ((arr1: Array, arr2: Array) => {
13 | const combinedMessages = [...arr1, ...arr2]
14 | const uniqueMessages = combinedMessages.filter((message, index, self) =>
15 | index === self.findIndex((p) => p.id === message.id)
16 | );
17 | return uniqueMessages.sort((a, b) => {
18 | if (!a.createdAt || !b.createdAt) throw new Error("createdAt is missing")
19 | return a.createdAt.localeCompare(b.createdAt)
20 | });
21 | })
22 |
23 | // Use ComponentProps to get the prop types of the Avatar component
24 | export type AuthorAvatarProps = {
25 | type: 'user' | 'gen-ai';
26 | name: string;
27 | initials?: string;
28 | loading?: boolean;
29 | } & Partial>;
30 | type AuthorsType = {
31 | [key: string]: AuthorAvatarProps;
32 | };
33 | export const AUTHORS: AuthorsType = {
34 | 'human': { type: 'user', name: 'Jane Doe', initials: 'JD' },
35 | 'ai': { type: 'gen-ai', name: 'Generative AI assistant' },
36 | 'tool': { type: 'gen-ai', name: 'Generative AI assistant' },
37 | };
38 | import Avatar from '@cloudscape-design/chat-components/avatar';
39 |
40 |
41 | export function ChatBubbleAvatar({ type, name, loading }: AuthorAvatarProps) {
42 | if (type === 'gen-ai') {
43 | return ;
44 | }
45 |
46 | return ;
47 | }
48 |
--------------------------------------------------------------------------------
/tailwind.config.ts:
--------------------------------------------------------------------------------
1 | import type { Config } from "tailwindcss";
2 |
3 | const config: Config = {
4 | content: [
5 | "./src/pages/**/*.{js,ts,jsx,tsx,mdx}",
6 | "./src/components/**/*.{js,ts,jsx,tsx,mdx}",
7 | "./src/app/**/*.{js,ts,jsx,tsx,mdx}",
8 | ],
9 | theme: {
10 | extend: {
11 | colors: {
12 | background: "var(--background)",
13 | foreground: "var(--foreground)",
14 | },
15 | },
16 | },
17 | plugins: [
18 | require('@tailwindcss/typography')
19 | ],
20 | };
21 | export default config;
22 |
--------------------------------------------------------------------------------
/test/integration/testAddIamDirective.ts:
--------------------------------------------------------------------------------
1 | import { handler } from "@/../amplify/functions/addIamDirectiveToAllAssets"
2 | import { AppSyncResolverEvent, Context } from 'aws-lambda';
3 | import { Schema } from '@/../amplify/data/resource';
4 | import outputs from '@/../amplify_outputs.json';
5 |
6 | // process.env.AMPLIFY_DATA_GRAPHQL_ENDPOINT = outputs.data.url
7 | process.env.ROOT_STACK_NAME = outputs.custom.root_stack_name
8 |
9 | const testArguments = {}
10 |
11 | const dummyContext: Context = {
12 | callbackWaitsForEmptyEventLoop: true,
13 | functionName: 'test-function',
14 | functionVersion: '$LATEST',
15 | invokedFunctionArn: 'arn:aws:lambda:us-east-1:123456789012:function:test-function',
16 | memoryLimitInMB: '128',
17 | awsRequestId: '52fdfc07-2182-154f-163f-5f0f9a621d72',
18 | logGroupName: '/aws/lambda/test-function',
19 | logStreamName: '2020/09/22/[$LATEST]abcdefghijklmnopqrstuvwxyz',
20 | // identity: null,
21 | // clientContext: null,
22 | getRemainingTimeInMillis: () => 3000,
23 | done: () => { },
24 | fail: () => { },
25 | succeed: () => { },
26 | };
27 |
28 | const main = async () => {
29 | const response = await handler({}, dummyContext, () => null)
30 |
31 | console.log('Handler response: ', response)
32 | }
33 |
34 | main()
--------------------------------------------------------------------------------
/test/integration/testApi.ts:
--------------------------------------------------------------------------------
1 | // import { generateAmplifyClientWrapper } from '../../amplify/functions/utils/amplifyUtils'
2 | // import outputs from '@/../amplify_outputs.json';
3 | // // import { convertPdfToImages } from '../../amplify/functions/graphql/queries'
4 |
5 | // // import { AppSyncClient } from '@aws-sdk/client-appsync';
6 |
7 | // // import AWS from 'aws-sdk';
8 |
9 | // // console.log('AWS Region: ', AWS.config.region)
10 | // // const credentials = new AWS.EnvironmentCredentials('AWS');
11 |
12 |
13 | // // console.log('amplifyClient: ', amplifyClientWrapper)
14 |
15 | // const main = async function () {
16 | // const creds = await credentials.getPromise();
17 | // console.log('creds: ', creds)
18 | // const env = {
19 | // AMPLIFY_DATA_GRAPHQL_ENDPOINT: outputs.data.url,
20 | // AWS_REGION: outputs.data.aws_region
21 | // }
22 |
23 | // console.log('env: ', env)
24 |
25 | // const amplifyClientWrapper = generateAmplifyClientWrapper(env)
26 |
27 | // const messages = await amplifyClientWrapper.getChatMessageHistory({
28 | // chatSessionId: "133dfbbb-359c-460a-989e-5d94b300da30",
29 | // latestHumanMessageText: "Hello World"
30 | // })
31 |
32 | // console.log("messages: ", messages)
33 |
34 |
35 | // // const convertPdfToImagesResponse = await amplifyClientWrapper.amplifyClient.graphql({
36 | // // query: convertPdfToImages,
37 | // // variables: {
38 | // // s3Key: "production-agent/well-files/field=SanJuanEast/uwi=30-039-07715/30-039-07715_00131.pdf"
39 | // // }
40 | // // })
41 | // // console.log('convertPdfToImagesResponse: ', convertPdfToImagesResponse)
42 |
43 | // }
44 |
45 | // main()
--------------------------------------------------------------------------------
/test/integration/testConfigureProdDb.ts:
--------------------------------------------------------------------------------
1 | import { handler } from "@/../amplify/functions/configureProdDb/index"
2 | import { Context } from 'aws-lambda';
3 | import outputs from '@/../amplify_outputs.json';
4 |
5 | import { getDeployedResourceArn, getLambdaEnvironmentVariables } from "../utils";
6 |
7 | const rootStackName = outputs.custom.root_stack_name
8 |
9 | const testArguments = {}
10 |
11 | const dummyContext: Context = {
12 | callbackWaitsForEmptyEventLoop: true,
13 | functionName: 'test-function',
14 | functionVersion: '$LATEST',
15 | invokedFunctionArn: 'arn:aws:lambda:us-east-1:123456789012:function:test-function',
16 | memoryLimitInMB: '128',
17 | awsRequestId: '52fdfc07-2182-154f-163f-5f0f9a621d72',
18 | logGroupName: '/aws/lambda/test-function',
19 | logStreamName: '2020/09/22/[$LATEST]abcdefghijklmnopqrstuvwxyz',
20 | // identity: null,
21 | // clientContext: null,
22 | getRemainingTimeInMillis: () => 3000,
23 | done: () => { },
24 | fail: () => { },
25 | succeed: () => { },
26 | };
27 |
28 | const main = async () => {
29 | process.env.ROOT_STACK_NAME = rootStackName
30 | // test()
31 | // getDeployedResourceArn(rootStackName, 'configureProdDbFunction')
32 | await getLambdaEnvironmentVariables(await getDeployedResourceArn(rootStackName, 'configureProdDbFunction'))
33 |
34 | const response = await handler({}, dummyContext, () => null)
35 |
36 | console.log('Handler response: ', response)
37 | }
38 |
39 | main()
--------------------------------------------------------------------------------
/test/integration/testConvertPdfToImageTool.ts:
--------------------------------------------------------------------------------
1 | // import {buildConvertPdfToImageTool} from '../../amplify/functions/productionAgentFunction/toolBox'
2 | // import outputs from '@/../amplify_outputs.json';
3 |
4 | // const convertpdfToImageTool = buildConvertPdfToImageTool({amplifyClient: outputs.storage.bucket_name})
5 |
6 | // async function main() {
7 | // const toolResponse = await convertpdfToImageTool.invoke({ s3Key: 'production-agent/well-files/field=SanJuanEast/uwi=30-039-07715/30-039-07715_00131.pdf'})
8 | // console.log(toolResponse)
9 | // }
10 |
11 | // main()
--------------------------------------------------------------------------------
/test/integration/testConvertPdfToYaml.ts:
--------------------------------------------------------------------------------
1 | import { handler } from "@/../amplify/functions/convertPdfToYaml/index"
2 | import { EventBridgeEvent, Context } from 'aws-lambda';
3 | import outputs from '@/../amplify_outputs.json';
4 |
5 | import { getDeployedResourceArn, getLambdaEnvironmentVariables } from "../utils";
6 |
7 | const rootStackName = outputs.custom.root_stack_name
8 |
9 |
10 | const dummyContext: Context = {
11 | callbackWaitsForEmptyEventLoop: true,
12 | functionName: 'test-function',
13 | functionVersion: '$LATEST',
14 | invokedFunctionArn: 'arn:aws:lambda:us-east-1:123456789012:function:test-function',
15 | memoryLimitInMB: '128',
16 | awsRequestId: '52fdfc07-2182-154f-163f-5f0f9a621d72',
17 | logGroupName: '/aws/lambda/test-function',
18 | logStreamName: '2020/09/22/[$LATEST]abcdefghijklmnopqrstuvwxyz',
19 | // identity: null,
20 | // clientContext: null,
21 | getRemainingTimeInMillis: () => 3000,
22 | done: () => { },
23 | fail: () => { },
24 | succeed: () => { },
25 | };
26 |
27 | const main = async () => {
28 | process.env.ROOT_STACK_NAME = rootStackName
29 | process.env.AWS_DEFAULT_REGION = outputs.auth.aws_region
30 |
31 | await getLambdaEnvironmentVariables(await getDeployedResourceArn(rootStackName, 'ConvertPdfToYamlFunction'))
32 |
33 | const testEvent = {
34 | "Records": [
35 | {
36 | "messageId": "7a984573-ccec-4e46-abea-1dd5ade5afbe",
37 | "receiptHandle": "AQEBQ9HJozlopJCagqLu17ajzLPgjB2+/BABuR1rSfzVQXvbYOPmZY5yiVAQaz0pyQj6hTeulKchxPZJ6abw3DucuZoobrcoM3msppVmZ76qPn8kqTMUJssW8wCTSJLwx97l+10k2qz/z1zDc3OH6m9W2Pea974Z2T50STtWHQY5KsSt+gGHKZjzoZflOldyh/6Qvh4gWXYpYMug8SJeNb6KzG0NQTUOAA+sUL6ckfYAEj8m/cgZ5i6lzU/hb4y7xEymL73GQWJLRWUDU8BsvllvX/EYboitrzIBTtURTG9gTh25chshV2nwjb0maTW5FDkY1wS/k/fxUcFnI5waOaaJC72gREnUMC9EnvZi22rSDG4ouWTY+Qqv8fjrdjJ6KAHt0oC2OBetvVZTsQYxHK4Ui49fxgayejnGeVijKcd96xNHMJ+YIZmwO+U9dMtm8Nict/Eb8npX/dIn5xKp+sEhXvkxALJ65PDtH6MQIG9mm4s=",
38 | "body": "{\"Records\":[{\"eventVersion\":\"2.1\",\"eventSource\":\"aws:s3\",\"awsRegion\":\"us-east-1\",\"eventTime\":\"2024-12-02T17:57:19.819Z\",\"eventName\":\"ObjectCreated:Copy\",\"userIdentity\":{\"principalId\":\"AWS:AROARQKFKHN2LJQXD5K3B:waltmayf+management@amazon.com\"},\"requestParameters\":{\"sourceIPAddress\":\"136.62.253.122\"},\"responseElements\":{\"x-amz-request-id\":\"Q1WGRGKKD9Q9SEHC\",\"x-amz-id-2\":\"ypCkHbztqIi4uVe+JRsxiOQKpDVHm7UyS5h/6S0IiEeU6p/467TmgLSMb7JgSvra1gUGUa3V5XncLs6KFOwemYpGH7CraNp3\"},\"s3\":{\"s3SchemaVersion\":\"1.0\",\"configurationId\":\"arn:aws:cloudformation:us-east-1:103761460084:stack/amplify-agentsforenergy-pw-sandbox-8d46451425-prodAgentStack8401ED6B-1X1UK0L60C7PU/68a16910-aced-11ef-b780-1263c9eff61d--6305615801843184238\",\"bucket\":{\"name\":\"amplify-agentsforenergy-pw-filedrivebucket01be03e1-oxivxspzy85y\",\"ownerIdentity\":{\"principalId\":\"A32XJER8F2MVSC\"},\"arn\":\"arn:aws:s3:::amplify-agentsforenergy-pw-filedrivebucket01be03e1-oxivxspzy85y\"},\"object\":{\"key\":\"production-agent/well-files/field%3DSanJuanEast/api%3D30-045-29202/30045292020000_12_wf.pdf\",\"size\":1103171,\"eTag\":\"58100e5bee694a518797e94c80f07559\",\"sequencer\":\"00674DF4FFA0F6ADFB\"}}}]}",
39 | "attributes": {
40 | "ApproximateReceiveCount": "3",
41 | "SentTimestamp": "1733162240572",
42 | "SenderId": "AROA4R74ZO52XAB5OD7T4:S3-PROD-END",
43 | "ApproximateFirstReceiveTimestamp": "1733162252325"
44 | },
45 | "messageAttributes": {},
46 | "md5OfBody": "d539ffe576e9720aa5147345a813c4ce",
47 | "eventSource": "aws:sqs",
48 | "eventSourceARN": "arn:aws:sqs:us-east-1:103761460084:amplify-agentsforenergy-pw-sandbox-8d4645142-PdfToYamlQueue4E2C37F2-MIsiQmYJE4vw",
49 | "awsRegion": "us-east-1"
50 | }
51 | ]
52 | }
53 |
54 | const response = await handler(testEvent, dummyContext, () => null)
55 |
56 | console.log('Handler response: ', response)
57 | }
58 |
59 | main()
--------------------------------------------------------------------------------
/test/integration/testGetPlanAndExecuteResponse.ts:
--------------------------------------------------------------------------------
1 | import { handler } from "@/../amplify/functions/planAndExecuteAgent/index"
2 | import { AppSyncResolverEvent, Context, AppSyncIdentity } from 'aws-lambda';
3 | import { Schema } from '@/../amplify/data/resource';
4 | import { STSClient } from "@aws-sdk/client-sts";
5 | import { AmplifyClientWrapper } from '@/../amplify/functions/utils/amplifyUtils'
6 | import { createChatSession } from "@/../amplify/functions/graphql/mutations";
7 |
8 | import { getDeployedResourceArn, getLambdaEnvironmentVariables } from "../utils";
9 |
10 | import outputs from '@/../amplify_outputs.json';
11 |
12 | const stsClient = new STSClient();
13 |
14 | const dummyContext: Context = {
15 | callbackWaitsForEmptyEventLoop: true,
16 | functionName: 'test-function',
17 | functionVersion: '$LATEST',
18 | invokedFunctionArn: 'arn:aws:lambda:us-east-1:123456789012:function:test-function',
19 | memoryLimitInMB: '128',
20 | awsRequestId: '52fdfc07-2182-154f-163f-5f0f9a621d72',
21 | logGroupName: '/aws/lambda/test-function',
22 | logStreamName: '2020/09/22/[$LATEST]abcdefghijklmnopqrstuvwxyz',
23 | // identity: null,
24 | // clientContext: null,
25 | getRemainingTimeInMillis: () => 3000,
26 | done: () => { },
27 | fail: () => { },
28 | succeed: () => { },
29 | };
30 |
31 | export const main = async () => {
32 | const rootStackName = outputs.custom.root_stack_name
33 | await getLambdaEnvironmentVariables(await getDeployedResourceArn(rootStackName, 'productionagentfunctionlambda'))
34 |
35 | process.env.AMPLIFY_DATA_GRAPHQL_ENDPOINT = outputs.data.url
36 | process.env.AWS_DEFAULT_REGION = outputs.auth.aws_region
37 | process.env.MODEL_ID = 'us.anthropic.claude-3-5-haiku-20241022-v1:0'
38 |
39 | const credentials = await stsClient.config.credentials()
40 | process.env.AWS_ACCESS_KEY_ID = credentials.accessKeyId
41 | process.env.AWS_SECRET_ACCESS_KEY = credentials.secretAccessKey
42 | process.env.AWS_SESSION_TOKEN = credentials.sessionToken
43 |
44 | const amplifyClientWrapper = new AmplifyClientWrapper({
45 | env: process.env
46 | })
47 |
48 | //Create a new chat session for testing
49 | const testChatSession = await amplifyClientWrapper.amplifyClient.graphql({ //To stream partial responces to the client
50 | query: createChatSession,
51 | variables: {
52 | input: {
53 | firstMessageSummary: "Test Summary"
54 | }
55 | }
56 | })
57 |
58 | const testArguments = {
59 | chatSessionId: testChatSession.data.createChatSession.id,
60 |
61 | // lastMessageText: `
62 | // What is the hometown of the 2015 Australian open winner?
63 | // `
64 |
65 | lastMessageText: `
66 | Execute a SQL query and plot the result to get the production over the last 12 weeks.
67 | Ust the ai role for all steps.
68 | `
69 | // "input": "What is 1+54?"
70 | }
71 |
72 | const event: AppSyncResolverEvent = {
73 | "arguments": testArguments,
74 | identity: { sub: "testIdentity" } as AppSyncIdentity,
75 | source: null,
76 | request: {
77 | headers: {},
78 | domainName: null,
79 | },
80 | info: {
81 | fieldName: 'yourFieldName',
82 | parentTypeName: 'Query',
83 | selectionSetList: [],
84 | selectionSetGraphQL: '',
85 | variables: {}
86 | },
87 | prev: null,
88 | stash: {},
89 | };
90 |
91 | const response = await handler(event, dummyContext, () => null)
92 |
93 | console.log('Handler response: ', response)
94 | }
95 |
96 | main()
--------------------------------------------------------------------------------
/test/integration/testGetProductionAgentResponse.ts:
--------------------------------------------------------------------------------
1 | import { handler } from "@/../amplify/functions/productionAgentFunction/index"
2 | import { AppSyncResolverEvent, Context, AppSyncIdentity } from 'aws-lambda';
3 | import { Schema } from '@/../amplify/data/resource';
4 | import { STSClient } from "@aws-sdk/client-sts";
5 | import { AmplifyClientWrapper } from '@/../amplify/functions/utils/amplifyUtils'
6 | import { createChatSession } from "@/../amplify/functions/graphql/mutations";
7 |
8 | import { getDeployedResourceArn, getLambdaEnvironmentVariables } from "../utils";
9 |
10 | import outputs from '@/../amplify_outputs.json';
11 |
12 | const stsClient = new STSClient();
13 |
14 | const dummyContext: Context = {
15 | callbackWaitsForEmptyEventLoop: true,
16 | functionName: 'test-function',
17 | functionVersion: '$LATEST',
18 | invokedFunctionArn: 'arn:aws:lambda:us-east-1:123456789012:function:test-function',
19 | memoryLimitInMB: '128',
20 | awsRequestId: '52fdfc07-2182-154f-163f-5f0f9a621d72',
21 | logGroupName: '/aws/lambda/test-function',
22 | logStreamName: '2020/09/22/[$LATEST]abcdefghijklmnopqrstuvwxyz',
23 | // identity: null,
24 | // clientContext: null,
25 | getRemainingTimeInMillis: () => 3000,
26 | done: () => { },
27 | fail: () => { },
28 | succeed: () => { },
29 | };
30 |
31 | export const main = async () => {
32 | const rootStackName = outputs.custom.root_stack_name
33 | await getLambdaEnvironmentVariables(await getDeployedResourceArn(rootStackName, 'productionagentfunctionlambda'))
34 |
35 | process.env.AMPLIFY_DATA_GRAPHQL_ENDPOINT = outputs.data.url
36 | process.env.AWS_DEFAULT_REGION = outputs.auth.aws_region
37 | process.env.MODEL_ID = 'us.anthropic.claude-3-sonnet-20240229-v1:0'
38 |
39 | const credentials = await stsClient.config.credentials()
40 | process.env.AWS_ACCESS_KEY_ID = credentials.accessKeyId
41 | process.env.AWS_SECRET_ACCESS_KEY = credentials.secretAccessKey
42 | process.env.AWS_SESSION_TOKEN = credentials.sessionToken
43 |
44 | const amplifyClientWrapper = new AmplifyClientWrapper({
45 | env: process.env
46 | })
47 |
48 | //Create a new chat session for testing
49 | const testChatSession = await amplifyClientWrapper.amplifyClient.graphql({ //To stream partial responces to the client
50 | query: createChatSession,
51 | variables: {
52 | input: {}
53 | }
54 | })
55 |
56 | const testArguments = {
57 | chatSessionId: testChatSession.data.createChatSession.id,
58 |
59 | // lastMessageText: `
60 | // What is the equation I should use for a decline curve analysis?
61 | // `
62 |
63 | lastMessageText: `
64 | For the well with API number 30-045-29202, estimate the Arps decline curve parameters.
65 | `
66 |
67 | // lastMessageText: `
68 | // Plot the total monthly oil, gas, and water production since 1900 for the well with API number 30-045-29202
69 | // `
70 | }
71 |
72 | const event: AppSyncResolverEvent = {
73 | "arguments": testArguments,
74 | identity: { sub: "testIdentity" } as AppSyncIdentity,
75 | source: null,
76 | request: {
77 | headers: {},
78 | domainName: null,
79 | },
80 | info: {
81 | fieldName: 'yourFieldName',
82 | parentTypeName: 'Query',
83 | selectionSetList: [],
84 | selectionSetGraphQL: '',
85 | variables: {}
86 | },
87 | prev: null,
88 | stash: {},
89 | };
90 |
91 | const response = await handler(event, dummyContext, () => null)
92 |
93 | console.log('Handler response: ', response)
94 | }
95 |
96 | main()
--------------------------------------------------------------------------------
/test/integration/testGetStructuredOutputRespose.ts:
--------------------------------------------------------------------------------
1 | import { handler } from "@/../amplify/functions/getStructuredOutputFromLangchain"
2 | import { AppSyncResolverEvent, Context, AppSyncIdentity } from 'aws-lambda';
3 | import { Schema } from '@/../amplify/data/resource';
4 | import { STSClient } from "@aws-sdk/client-sts";
5 | import { AmplifyClientWrapper } from '@/../amplify/functions/utils/amplifyUtils'
6 | import { createChatSession } from "@/../amplify/functions/graphql/mutations";
7 |
8 | import { getDeployedResourceArn, getLambdaEnvironmentVariables } from "../utils";
9 |
10 | import outputs from '@/../amplify_outputs.json';
11 |
12 | const stsClient = new STSClient();
13 |
14 | const dummyContext: Context = {
15 | callbackWaitsForEmptyEventLoop: true,
16 | functionName: 'test-function',
17 | functionVersion: '$LATEST',
18 | invokedFunctionArn: 'arn:aws:lambda:us-east-1:123456789012:function:test-function',
19 | memoryLimitInMB: '128',
20 | awsRequestId: '52fdfc07-2182-154f-163f-5f0f9a621d72',
21 | logGroupName: '/aws/lambda/test-function',
22 | logStreamName: '2020/09/22/[$LATEST]abcdefghijklmnopqrstuvwxyz',
23 | // identity: null,
24 | // clientContext: null,
25 | getRemainingTimeInMillis: () => 3000,
26 | done: () => { },
27 | fail: () => { },
28 | succeed: () => { },
29 | };
30 |
31 | export const main = async () => {
32 | const rootStackName = outputs.custom.root_stack_name
33 | await getLambdaEnvironmentVariables(await getDeployedResourceArn(rootStackName, 'getstructuredoutputlambda'))
34 |
35 | process.env.AMPLIFY_DATA_GRAPHQL_ENDPOINT = outputs.data.url
36 | process.env.AWS_DEFAULT_REGION = outputs.auth.aws_region
37 | process.env.MODEL_ID = 'anthropic.claude-3-sonnet-20240229-v1:0'
38 |
39 | const credentials = await stsClient.config.credentials()
40 | process.env.AWS_ACCESS_KEY_ID = credentials.accessKeyId
41 | process.env.AWS_SECRET_ACCESS_KEY = credentials.secretAccessKey
42 | process.env.AWS_SESSION_TOKEN = credentials.sessionToken
43 |
44 | const amplifyClientWrapper = new AmplifyClientWrapper({
45 | env: process.env
46 | })
47 |
48 | //Create a new chat session for testing
49 | const testChatSession = await amplifyClientWrapper.amplifyClient.graphql({ //To stream partial responces to the client
50 | query: createChatSession,
51 | variables: {
52 | input: {}
53 | }
54 | })
55 |
56 | const testArguments = {
57 | chatSessionId: testChatSession.data.createChatSession.id,
58 | lastMessageText: 'hello',
59 | outputStructure: JSON.stringify({
60 | type: 'object',
61 | properties: {
62 | name: {
63 | type: 'string',
64 | description: 'The name of the person'
65 | },
66 | age: {
67 | type: 'number',
68 | description: 'The age of the person'
69 | }
70 | },
71 | required: ['name', 'age']
72 | })
73 |
74 | }
75 |
76 | const event: AppSyncResolverEvent = {
77 | "arguments": testArguments,
78 | identity: { sub: "testIdentity" } as AppSyncIdentity,
79 | source: null,
80 | request: {
81 | headers: {},
82 | domainName: null,
83 | },
84 | info: {
85 | fieldName: 'yourFieldName',
86 | parentTypeName: 'Query',
87 | selectionSetList: [],
88 | selectionSetGraphQL: '',
89 | variables: {}
90 | },
91 | prev: null,
92 | stash: {},
93 | };
94 |
95 | const response = await handler(event, dummyContext, () => null)
96 |
97 | console.log('Handler response: ', response)
98 | }
99 |
100 | main()
--------------------------------------------------------------------------------
/test/integration/testModifyPlanAndExecute.ts:
--------------------------------------------------------------------------------
1 | import { handler } from "@/../amplify/functions/planAndExecuteAgent/index"
2 | import { AppSyncResolverEvent, Context, AppSyncIdentity } from 'aws-lambda';
3 | import { Schema } from '@/../amplify/data/resource';
4 | import { STSClient } from "@aws-sdk/client-sts";
5 | import { AmplifyClientWrapper } from '@/../amplify/functions/utils/amplifyUtils'
6 | import { createChatSession } from "@/../amplify/functions/graphql/mutations";
7 |
8 | import { getDeployedResourceArn, getLambdaEnvironmentVariables } from "../utils";
9 |
10 | import outputs from '@/../amplify_outputs.json';
11 |
12 | const stsClient = new STSClient();
13 |
14 | const dummyContext: Context = {
15 | callbackWaitsForEmptyEventLoop: true,
16 | functionName: 'test-function',
17 | functionVersion: '$LATEST',
18 | invokedFunctionArn: 'arn:aws:lambda:us-east-1:123456789012:function:test-function',
19 | memoryLimitInMB: '128',
20 | awsRequestId: '52fdfc07-2182-154f-163f-5f0f9a621d72',
21 | logGroupName: '/aws/lambda/test-function',
22 | logStreamName: '2020/09/22/[$LATEST]abcdefghijklmnopqrstuvwxyz',
23 | // identity: null,
24 | // clientContext: null,
25 | getRemainingTimeInMillis: () => 3000,
26 | done: () => { },
27 | fail: () => { },
28 | succeed: () => { },
29 | };
30 |
31 | export const main = async () => {
32 | const rootStackName = outputs.custom.root_stack_name
33 | await getLambdaEnvironmentVariables(await getDeployedResourceArn(rootStackName, 'productionagentfunctionlambda'))
34 |
35 | process.env.AMPLIFY_DATA_GRAPHQL_ENDPOINT = outputs.data.url
36 | process.env.AWS_DEFAULT_REGION = outputs.auth.aws_region
37 | process.env.MODEL_ID = 'us.anthropic.claude-3-5-haiku-20241022-v1:0'
38 |
39 | const credentials = await stsClient.config.credentials()
40 | process.env.AWS_ACCESS_KEY_ID = credentials.accessKeyId
41 | process.env.AWS_SECRET_ACCESS_KEY = credentials.secretAccessKey
42 | process.env.AWS_SESSION_TOKEN = credentials.sessionToken
43 |
44 | const amplifyClientWrapper = new AmplifyClientWrapper({
45 | env: process.env
46 | })
47 |
48 | //Create a new chat session for testing
49 | const testChatSession = await amplifyClientWrapper.amplifyClient.graphql({ //To stream partial responces to the client
50 | query: createChatSession,
51 | variables: {
52 | input: {
53 | firstMessageSummary: "Test Summary",
54 | planGoal: "Where should I go on vacation?",
55 | planSteps: [
56 | JSON.stringify({
57 | title: "Get user beach / snow preferance",
58 | description: "Get the user's beach / snow preferance",
59 | role: "human"
60 | })
61 | ]
62 | }
63 | }
64 | })
65 |
66 | const testArguments = {
67 | chatSessionId: testChatSession.data.createChatSession.id,
68 |
69 | lastMessageText: `
70 | I like scuba diving.
71 | `
72 | }
73 |
74 | const event: AppSyncResolverEvent = {
75 | "arguments": testArguments,
76 | identity: { sub: "testIdentity" } as AppSyncIdentity,
77 | source: null,
78 | request: {
79 | headers: {},
80 | domainName: null,
81 | },
82 | info: {
83 | fieldName: 'yourFieldName',
84 | parentTypeName: 'Query',
85 | selectionSetList: [],
86 | selectionSetGraphQL: '',
87 | variables: {}
88 | },
89 | prev: null,
90 | stash: {},
91 | };
92 |
93 | const response = await handler(event, dummyContext, () => null)
94 |
95 | console.log('Handler response: ', response)
96 | }
97 |
98 | main()
--------------------------------------------------------------------------------
/test/unit/testAthenaFederatedQuery.ts:
--------------------------------------------------------------------------------
1 | import { getDeployedResourceArn, getLambdaEnvironmentVariables } from "../utils";
2 | import { executeSQLQueryTool } from '../../amplify/functions/productionAgentFunction/toolBox';
3 | import outputs from '@/../amplify_outputs.json';
4 |
5 | async function main() {
6 | const rootStackName = outputs.custom.root_stack_name
7 | const sampleAthenaDataSource = await getDeployedResourceArn(rootStackName, `PostgresAthenaDataSource`)
8 | console.log('sampleAthenaDataSource: ', sampleAthenaDataSource)
9 |
10 | await getLambdaEnvironmentVariables(await getDeployedResourceArn(rootStackName, 'productionagentfunctionlambda'))
11 |
12 | // console.log('ATHENA_WORKGROUP_NAME: ', process.env.ATHENA_WORKGROUP_NAME)
13 |
14 | const tableDefinitions = await executeSQLQueryTool.invoke({
15 | query: /* sql */ `
16 | SELECT
17 | oil ,
18 | gas ,
19 | water,
20 | proddate
21 | FROM "${sampleAthenaDataSource}".production.daily
22 | WHERE proddate >= date_add('week', -12, current_date)`,
23 |
24 |
25 | // query: /* sql */ `
26 | // SHOW TABLES FROM AwsDataCatalog.prod_db_e1d;
27 | // `,
28 |
29 | // query: /* sql */ `
30 | // SELECT schema_name
31 | // FROM information_schema.schemata;
32 | // `,
33 | // database: "public",
34 | // columnNameFromQueryForXAxis: 'proddate',
35 | // chartTitle: "test chart"
36 | });
37 | console.log('result:\n', tableDefinitions);
38 | }
39 |
40 | main()
--------------------------------------------------------------------------------
/test/unit/testBedrockKnowlegeBaseRetriever.ts:
--------------------------------------------------------------------------------
1 | import { getDeployedResourceArn, getLambdaEnvironmentVariables } from "../utils";
2 | import { getTableDefinitionsTool } from '../../amplify/functions/productionAgentFunction/toolBox';
3 | import outputs from '@/../amplify_outputs.json';
4 |
5 | async function main() {
6 | const rootStackName = outputs.custom.root_stack_name
7 | await getLambdaEnvironmentVariables(await getDeployedResourceArn(rootStackName, 'productionagentfunctionlambda'))
8 |
9 | const tableDefinitions = await getTableDefinitionsTool.invoke({ tableFeatures: "Average Oil Production Over the Last 10 weeks" });
10 | console.log('tableDefinitions:\n', tableDefinitions);
11 | }
12 |
13 | main()
--------------------------------------------------------------------------------
/test/unit/testGetStructuredOutput.ts:
--------------------------------------------------------------------------------
1 | import { getStructuredOutputResponse } from '@/../amplify/functions/getStructuredOutputFromLangchain'
2 | import { HumanMessage } from "@langchain/core/messages";
3 | import outputs from '@/../amplify_outputs.json';
4 |
5 | const main = async () => {
6 | process.env.AWS_DEFAULT_REGION = outputs.auth.aws_region
7 |
8 | const outputStructure = {
9 | title: "SummarizeMessageIntnet",
10 | description: "Summarize the intent of the user's message?",
11 | type: "object",
12 | properties: {
13 | summary: {
14 | type: 'string',
15 | description: `Message intent summary in 20 characters or fewer.`,
16 | },
17 | svgImage: {
18 | type: 'string',
19 | description: `SVG describing the input`,
20 | }
21 | },
22 | required: ['summary','svgImage'],
23 | };
24 |
25 | const response = await getStructuredOutputResponse({
26 | messages: [
27 | new HumanMessage({ content: "I'm the strongest, greatest, most hansome man in the world" })
28 | ],
29 | modelId: 'us.anthropic.claude-3-sonnet-20240229-v1:0',
30 | outputStructure: outputStructure
31 | })
32 | console.log(response)
33 | }
34 |
35 | main()
--------------------------------------------------------------------------------
/test/unit/testInvokeBedrockAgent.ts:
--------------------------------------------------------------------------------
1 | import { handler } from "@/../amplify/functions/invokeBedrockAgent"
2 | import { AppSyncResolverEvent, Context, AppSyncIdentity } from 'aws-lambda';
3 | import { Schema } from '@/../amplify/data/resource';
4 | import { STSClient } from "@aws-sdk/client-sts";
5 | import { AmplifyClientWrapper, createChatMessage } from '@/../amplify/functions/utils/amplifyUtils'
6 | import { ChatMessageRole } from '@/../amplify/functions/graphql/API'
7 | import { createChatSession } from "@/../amplify/functions/graphql/mutations";
8 |
9 | import { getDeployedResourceArn, getLambdaEnvironmentVariables } from "../utils";
10 |
11 | import outputs from '@/../amplify_outputs.json';
12 |
13 | const stsClient = new STSClient();
14 |
15 | const dummyContext: Context = {
16 | callbackWaitsForEmptyEventLoop: true,
17 | functionName: 'test-function',
18 | functionVersion: '$LATEST',
19 | invokedFunctionArn: 'arn:aws:lambda:us-east-1:123456789012:function:test-function',
20 | memoryLimitInMB: '128',
21 | awsRequestId: '52fdfc07-2182-154f-163f-5f0f9a621d72',
22 | logGroupName: '/aws/lambda/test-function',
23 | logStreamName: '2020/09/22/[$LATEST]abcdefghijklmnopqrstuvwxyz',
24 | // identity: null,
25 | // clientContext: null,
26 | getRemainingTimeInMillis: () => 3000,
27 | done: () => { },
28 | fail: () => { },
29 | succeed: () => { },
30 | };
31 |
32 | export const main = async () => {
33 | const rootStackName = outputs.custom.root_stack_name
34 | await getLambdaEnvironmentVariables(await getDeployedResourceArn(rootStackName, 'invokebedrockagentlambda'))
35 |
36 | process.env.AMPLIFY_DATA_GRAPHQL_ENDPOINT = outputs.data.url
37 | process.env.AWS_DEFAULT_REGION = outputs.auth.aws_region
38 | process.env.MODEL_ID = 'anthropic.claude-3-sonnet-20240229-v1:0'
39 |
40 | const credentials = await stsClient.config.credentials()
41 | process.env.AWS_ACCESS_KEY_ID = credentials.accessKeyId
42 | process.env.AWS_SECRET_ACCESS_KEY = credentials.secretAccessKey
43 | process.env.AWS_SESSION_TOKEN = credentials.sessionToken
44 |
45 | const amplifyClientWrapper = new AmplifyClientWrapper({
46 | env: process.env
47 | })
48 |
49 | //Create a new chat session for testing
50 | const testChatSession = await amplifyClientWrapper.amplifyClient.graphql({ //To stream partial responces to the client
51 | query: createChatSession,
52 | variables: {
53 | input: {}
54 | }
55 | })
56 |
57 | const testChatMessage = await amplifyClientWrapper.amplifyClient.graphql({
58 | query: createChatMessage,
59 | variables: {
60 | input: {
61 | chatSessionId: testChatSession.data.createChatSession.id,
62 | content: 'How many tanks are in my biodiesel unit?',
63 | role: ChatMessageRole.human,
64 | trace: "Test Trace"
65 | }
66 | }
67 | })
68 |
69 | console.log("test message:\n", testChatMessage.data.createChatMessage)
70 |
71 | const testArguments = {
72 | agentAliasId: outputs.custom.maintenanceAgentAliasId,
73 | agentId: outputs.custom.maintenanceAgentId,
74 | chatSessionId: testChatSession.data.createChatSession.id,
75 | prompt: 'How many tanks are in my biodiesel unit?',
76 | }
77 |
78 | const event: AppSyncResolverEvent = {
79 | arguments: testArguments,
80 | identity: { sub: "testIdentity" } as AppSyncIdentity,
81 | source: null,
82 | request: {
83 | headers: {},
84 | domainName: null,
85 | },
86 | info: {
87 | fieldName: 'yourFieldName',
88 | parentTypeName: 'Query',
89 | selectionSetList: [],
90 | selectionSetGraphQL: '',
91 | variables: {}
92 | },
93 | prev: null,
94 | stash: {},
95 | };
96 |
97 | const response = await handler(event, dummyContext, () => null)
98 |
99 | console.log('Handler response: ', response)
100 | }
101 |
102 | main()
--------------------------------------------------------------------------------
/test/unit/testWellTableTool.ts:
--------------------------------------------------------------------------------
1 | import { STSClient } from "@aws-sdk/client-sts";
2 | import { z } from 'zod';
3 |
4 | import { getDeployedResourceArn, getLambdaEnvironmentVariables } from "../utils";
5 | import { wellTableTool, wellTableSchema } from '../../amplify/functions/productionAgentFunction/toolBox';
6 | import { AmplifyClientWrapper } from '@/../amplify/functions/utils/amplifyUtils'
7 | import outputs from '@/../amplify_outputs.json';
8 |
9 | const testArguments =
10 | {
11 | "wellApiNumber": "30-045-29202",
12 | "tableColumns": [
13 | {
14 | "columnDescription": "The type of operation performed on the well",
15 | "columnName": "Operation Type",
16 | "columnDataDefinition": {
17 | "type": "string",
18 | "enum": [
19 | "drill",
20 | "completion",
21 | "transportation",
22 | "cathodic protection"
23 | ]
24 | }
25 | },
26 | {
27 | "columnDescription": "Text describing the details of the operation",
28 | "columnName": "Operation Details",
29 | "columnDataDefinition": {
30 | "type": "string",
31 | }
32 | }
33 | ],
34 | // "dataToExclude": "transportation corporation, cathotic protection",
35 | // s3Key: "production-agent/well-files/field=SanJuanEast/uwi=30-039-07715/30-039-07715_00114.pdf" // Change in Transporter
36 | // s3Key: "production-agent/well-files/field=SanJuanEast/uwi=30-039-07715/3003907715_24_wf_1.pdf" // Cathodic Protection
37 | // "s3Key": "production-agent/well-files/field=SanJuanEast/uwi=30-039-07715/30-039-07715_00112.pdf" //Drill Report
38 | // s3Key: "production-agent/well-files/field=SanJuanEast/uwi=30-039-07715/30-039-07715_00117.pdf" // Drill Report
39 | "s3Key": "production-agent/well-files/field=SanJuanEast/uwi=30-039-07715/30-039-07715_00131.pdf" // Drill Report
40 | } as z.infer
41 |
42 | async function main() {
43 | const rootStackName = outputs.custom.root_stack_name
44 | await getLambdaEnvironmentVariables(await getDeployedResourceArn(rootStackName, 'productionagentfunctionlambda'))
45 | process.env.AMPLIFY_DATA_GRAPHQL_ENDPOINT = outputs.data.url
46 | process.env.AWS_DEFAULT_REGION = outputs.auth.aws_region
47 | process.env.MODEL_ID = 'us.anthropic.claude-3-sonnet-20240229-v1:0'
48 |
49 | const tableDefinitions = await wellTableTool.invoke(testArguments);
50 | console.log('tableDefinitions:\n', tableDefinitions);
51 | }
52 |
53 | main()
--------------------------------------------------------------------------------
/test/utils.ts:
--------------------------------------------------------------------------------
1 | import {
2 | CloudFormationClient,
3 | ListStackResourcesCommand,
4 | } from "@aws-sdk/client-cloudformation"
5 | import {
6 | LambdaClient,
7 | GetFunctionConfigurationCommand
8 | } from "@aws-sdk/client-lambda";
9 |
10 | // export async function test() {
11 | // console.log('test')
12 | // }
13 |
14 | export async function getDeployedResourceArn(
15 | rootStackName: string,
16 | targetLogicalIdPrefix: string
17 | ): Promise {
18 | const cloudformation = new CloudFormationClient();
19 |
20 | async function searchStack(stackName: string): Promise {
21 | try {
22 | const resources = await cloudformation.send(new ListStackResourcesCommand({
23 | StackName: stackName,
24 | }))
25 |
26 | if (!resources || !resources.StackResourceSummaries) throw new Error(`No resources found in stack ${stackName}`);
27 |
28 | for (const resource of resources.StackResourceSummaries || []) {
29 | if (resource && resource.LogicalResourceId &&
30 | (
31 | resource.LogicalResourceId.slice(0,-8) === targetLogicalIdPrefix ||
32 | resource.LogicalResourceId === targetLogicalIdPrefix
33 | )
34 | ) {
35 | return resource.PhysicalResourceId;
36 | }
37 |
38 | if (resource.ResourceType === 'AWS::CloudFormation::Stack') {
39 | const nestedStackArn = resource.PhysicalResourceId;
40 | if (nestedStackArn) {
41 | const result = await searchStack(nestedStackArn);
42 | if (result) return result;
43 | }
44 | }
45 | }
46 |
47 | // If we've gone through all resources and haven't returned, check if there's a next token
48 | if (resources.NextToken) {
49 | const nextResources = await cloudformation.send(new ListStackResourcesCommand({
50 | StackName: stackName,
51 | }))
52 | resources.StackResourceSummaries?.push(...(nextResources.StackResourceSummaries || []));
53 | }
54 |
55 | } catch (error) {
56 | console.error(`Error searching stack ${stackName}:`, error);
57 | }
58 |
59 | return undefined;
60 | }
61 |
62 | const resourceId = await searchStack(rootStackName)
63 | if (!resourceId) throw new Error(`Could not find resource with logical ID: ${targetLogicalIdPrefix}`);
64 |
65 | console.log(`For logical id ${targetLogicalIdPrefix}, found PhysicalResourceId ${resourceId}`)
66 | return resourceId;
67 | }
68 |
69 |
70 | export async function getLambdaEnvironmentVariables(functionName: string): Promise {
71 | try {
72 | // Initialize the Lambda client
73 | const client = new LambdaClient();
74 |
75 | // Create the command to get function configuration
76 | const command = new GetFunctionConfigurationCommand({
77 | FunctionName: functionName
78 | });
79 |
80 | // Get the function configuration
81 | const response = await client.send(command);
82 |
83 | // Check if environment variables exist
84 | if (response.Environment && response.Environment.Variables) {
85 | const envVars = response.Environment.Variables;
86 |
87 | // Set each environment variable locally
88 | for (const [key, value] of Object.entries(envVars)) {
89 | if (value) {
90 | process.env[key] = value;
91 | console.log(`Set ${key} environment variable to ${value}`);
92 | }
93 | }
94 | } else {
95 | console.log('No environment variables found for the specified Lambda function');
96 | }
97 |
98 | } catch (error) {
99 | console.error('Error retrieving Lambda environment variables:', error);
100 | throw error;
101 | }
102 | }
--------------------------------------------------------------------------------
/tsconfig.json:
--------------------------------------------------------------------------------
1 | {
2 | "compilerOptions": {
3 | "target": "es2015",
4 | "lib": ["dom", "dom.iterable", "esnext"],
5 | "allowJs": true,
6 | "skipLibCheck": true,
7 | "strict": true,
8 | "noEmit": true,
9 | "esModuleInterop": true,
10 | "module": "esnext",
11 | "moduleResolution": "bundler",
12 | "resolveJsonModule": true,
13 | "isolatedModules": true,
14 | "jsx": "preserve",
15 | "incremental": true,
16 | "plugins": [
17 | {
18 | "name": "next"
19 | }
20 | ],
21 | "paths": {
22 | "@/*": ["./src/*"],
23 | "$amplify/*": [
24 | "./.amplify/generated/*"
25 | ],
26 | "$amplify/env/*": ["./.amplify/generated/env/*"]
27 | }
28 | },
29 | "include": ["next-env.d.ts", "**/*.ts", "**/*.tsx", ".next/types/**/*.ts"],
30 | "exclude": ["node_modules", "amplify/**/*", "tmp/**/*"],
31 | "references": [
32 | {
33 | "path": "./amplify"
34 | }
35 | ]
36 | }
37 |
--------------------------------------------------------------------------------