-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
add document check node, update variable parsing in llm completion node
- Loading branch information
Showing
11 changed files
with
297 additions
and
35 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
67 changes: 67 additions & 0 deletions
67
packages/@pufflig/ps-nodes-config/src/modifiers/document_check/document_check.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,67 @@ | ||
| import { chat_models, completion_models, default_completion_model } from "@pufflig/ps-models"; | ||
| import { NodeConfig } from "@pufflig/ps-types"; | ||
|
|
||
| export const documentCheckNodeType = "modifier/document_check" as const; | ||
|
|
||
| export const documentCheck: NodeConfig = { | ||
| name: "Document Check", | ||
| description: "Run a checklist or extract information from a document.", | ||
| tags: ["modifier", "document", "text"], | ||
| status: "stable", | ||
| execution: { | ||
| inputs: [ | ||
| { | ||
| id: "exec:input", | ||
| }, | ||
| ], | ||
| outputs: [ | ||
| { | ||
| id: "exec:output", | ||
| name: "Completed", | ||
| }, | ||
| ], | ||
| }, | ||
| outputs: [ | ||
| { | ||
| id: "list", | ||
| name: "List", | ||
| description: "A list, checklist or other information about the document", | ||
| type: "text", | ||
| defaultValue: "", | ||
| }, | ||
| ], | ||
| inputs: [ | ||
| { | ||
| id: "model", | ||
| name: "Model", | ||
| description: "The model to use", | ||
| type: "model", | ||
| definition: { ...completion_models, ...chat_models }, | ||
| defaultValue: { | ||
| modelId: default_completion_model, | ||
| parameters: {}, | ||
| }, | ||
| }, | ||
| { | ||
| id: "prompt", | ||
| name: "Prompt", | ||
| description: "Prompt to check the document with", | ||
| type: "text", | ||
| defaultValue: `Extract information in the document below and insert them in the csv table, don't overwrite existing values and keep things empty if you cannot find information in the document:\n\nTABLE EXAMPLE:\ncharacters, age\nmickey mouse, 10\ndonald duck, -\n\nTABLE:\n[[table]]\n\nDOCUMENT:\n[[document]]\n\nTABLE:\n`, | ||
| }, | ||
| { | ||
| id: "table", | ||
| name: "Table", | ||
| description: "The list, table or checklist to parse the document with.", | ||
| type: "text", | ||
| defaultValue: "", | ||
| }, | ||
| { | ||
| id: "document", | ||
| name: "Document", | ||
| description: "Document to be processed", | ||
| type: "text", | ||
| defaultValue: "", | ||
| }, | ||
| ], | ||
| }; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
91 changes: 91 additions & 0 deletions
91
packages/@pufflig/ps-nodes/src/modifiers/document_check/document_check.test.ts
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,91 @@ | ||
| import { execute, LLMCompletionInput } from "./document_check"; | ||
| import axios from "axios"; | ||
|
|
||
| jest.mock("axios"); | ||
|
|
||
| describe("documentCheck", () => { | ||
| beforeEach(() => { | ||
| jest.resetAllMocks(); | ||
| }); | ||
|
|
||
| it("should return the completion string", async () => { | ||
| const input: LLMCompletionInput = { | ||
| prompt: "Hello, world!", | ||
| model: { | ||
| modelId: "test_model", | ||
| parameters: {}, | ||
| }, | ||
| document: "This is a test document.", | ||
| table: "test_table", | ||
| }; | ||
|
|
||
| const expectedOutput = { result: "This is a test completion." }; | ||
| const mockedAxiosResponse = { data: expectedOutput }; | ||
| (axios.post as jest.MockedFunction<typeof axios.post>).mockResolvedValueOnce(mockedAxiosResponse); | ||
|
|
||
| const output = await execute(input); | ||
|
|
||
| expect(output).toEqual({ completion: "This is a test completion." }); | ||
| expect(axios.post).toHaveBeenCalledTimes(1); | ||
| }); | ||
|
|
||
| it("should parse input variables", async () => { | ||
| const input: LLMCompletionInput = { | ||
| prompt: "Hello, {{myVariable}}!", | ||
| model: { | ||
| modelId: "test_model", | ||
| parameters: {}, | ||
| }, | ||
| document: "This is a test document.", | ||
| table: "test_table", | ||
| myVariable: "myValue", | ||
| }; | ||
|
|
||
| const expectedOutput = { result: "This is a test completion." }; | ||
| const mockedAxiosResponse = { data: expectedOutput }; | ||
| (axios.post as jest.MockedFunction<typeof axios.post>).mockResolvedValueOnce(mockedAxiosResponse); | ||
|
|
||
| const output = await execute(input); | ||
|
|
||
| expect(output).toEqual({ completion: "This is a test completion." }); | ||
| expect(axios.post).toHaveBeenCalledTimes(1); | ||
| expect(axios.post).toHaveBeenCalledWith( | ||
| expect.any(String), | ||
| { | ||
| document: "This is a test document.", | ||
| format: "test_table", | ||
| modelId: "test_model", | ||
| options: { | ||
| cache: true, | ||
| track: true, | ||
| }, | ||
| parameters: {}, | ||
| prompt: "Hello, myValue!", | ||
| }, | ||
| { | ||
| headers: { | ||
| Authorization: "Bearer undefined", | ||
| "Content-Type": "application/json", | ||
| }, | ||
| } | ||
| ); | ||
| }); | ||
|
|
||
| it("should throw an error if the API call fails", async () => { | ||
| const input: LLMCompletionInput = { | ||
| prompt: "Hello, world!", | ||
| model: { | ||
| modelId: "test_model", | ||
| parameters: {}, | ||
| }, | ||
| document: "This is a test document.", | ||
| table: "test_table", | ||
| }; | ||
|
|
||
| const expectedError = new Error("API call failed."); | ||
| (axios.post as jest.MockedFunction<typeof axios.post>).mockRejectedValueOnce(expectedError); | ||
|
|
||
| await expect(execute(input)).rejects.toThrow(expectedError); | ||
| expect(axios.post).toHaveBeenCalledTimes(1); | ||
| }); | ||
| }); |
Oops, something went wrong.