fix: add retriever input_type on RetrieverToolComponent + tests (#3601)

*  (RetrieverTool.py): add "input_types" field to the configuration of the RetrieverToolComponent class

*  (generalBugs-shard-12.spec.ts): add end-to-end test for connecting RetrieverTool to other components in the application.
This commit is contained in:
Cristhian Zanforlin Lousa 2024-08-28 16:02:42 -03:00 committed by GitHub
commit 7641e81ca3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 98 additions and 0 deletions

View file

@ -15,6 +15,7 @@ class RetrieverToolComponent(CustomComponent):
"display_name": "Retriever",
"info": "Retriever to interact with",
"type": BaseRetriever,
"input_types": ["Retriever"],
},
"name": {"display_name": "Name", "info": "Name of the tool"},
"description": {"display_name": "Description", "info": "Description of the tool"},

View file

@ -0,0 +1,97 @@
import { expect, test } from "@playwright/test";
test("user should be able to connect RetrieverTool to another components", async ({
page,
}) => {
await page.goto("/");
await page.waitForSelector('[data-testid="mainpage_title"]', {
timeout: 30000,
});
await page.waitForSelector('[id="new-project-btn"]', {
timeout: 30000,
});
let modalCount = 0;
try {
const modalTitleElement = await page.getByTestId("modal-title");
modalCount = await modalTitleElement.count();
} catch (error) {
modalCount = 0;
}
while (modalCount === 0) {
await page.getByText("New Project", { exact: true }).click();
await page.waitForTimeout(3000);
modalCount = await page.getByTestId("modal-title").count();
}
await page.waitForSelector('[data-testid="blank-flow"]', {
timeout: 30000,
});
await page.getByTestId("blank-flow").click();
await page.waitForSelector('[data-testid="extended-disclosure"]', {
timeout: 30000,
});
await page.getByTestId("extended-disclosure").click();
await page.getByPlaceholder("Search").click();
await page.getByPlaceholder("Search").fill("retriever");
await page.waitForTimeout(1000);
let modelElement = await page.getByTestId("toolsRetrieverTool");
let targetElement = await page.locator('//*[@id="react-flow-id"]');
await modelElement.dragTo(targetElement);
await page.mouse.up();
await page.mouse.down();
await page.getByTitle("fit view").click();
await page.getByTitle("zoom out").click();
await page.waitForTimeout(1000);
await page.getByTitle("zoom out").click();
await page
.locator('//*[@id="react-flow-id"]')
.hover()
.then(async () => {
await page.mouse.down();
await page.mouse.move(-500, 500);
});
await page.mouse.up();
await page.getByPlaceholder("Search").click();
await page.getByPlaceholder("Search").fill("chroma");
await page.waitForTimeout(1000);
modelElement = await page.getByTestId("vectorstoresChroma DB");
targetElement = await page.locator('//*[@id="react-flow-id"]');
await modelElement.dragTo(targetElement);
await page.mouse.up();
await page.mouse.down();
await page.waitForTimeout(1000);
await page.getByTitle("fit view").click();
await page.getByTitle("fit view").click();
//connection
const chromaDbOutput = await page
.getByTestId("handle-chroma-shownode-retriever-right")
.nth(0);
await chromaDbOutput.hover();
await page.mouse.down();
const retrieverToolInput = await page
.getByTestId("handle-retrievertool-shownode-retriever-left")
.nth(0);
await retrieverToolInput.hover();
await page.mouse.up();
await page.waitForTimeout(1000);
expect(await page.locator(".react-flow__edge-interaction").count()).toBe(1);
});