From 2860e0e59188dc33161170c0e434a69d6580c64d Mon Sep 17 00:00:00 2001 From: Viktor Avelino <64113566+viktoravelino@users.noreply.github.com> Date: Mon, 27 Apr 2026 15:19:14 -0400 Subject: [PATCH 1/3] fix(frontend): validate deployment agent name (#12896) --- .../__tests__/create-mode.test.tsx | 36 +++++++++++++++++++ .../__tests__/custom-tool-naming.test.tsx | 14 ++++++++ .../__tests__/edit-mode.test.tsx | 25 +++++++++++++ .../__tests__/step-type.test.tsx | 30 ++++++++++++++++ .../deploymentsPage/components/step-type.tsx | 13 ++++++- .../contexts/deployment-stepper-context.tsx | 26 ++++++++++++-- 6 files changed, 140 insertions(+), 4 deletions(-) diff --git a/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/create-mode.test.tsx b/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/create-mode.test.tsx index c0d7f1e73f..074de71ac4 100644 --- a/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/create-mode.test.tsx +++ b/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/create-mode.test.tsx @@ -218,6 +218,42 @@ describe("Create mode — canGoNext validation", () => { expect(result.current.canGoNext).toBe(true); }); + + it("blocks when trimmed name does not start with a letter", () => { + const { result } = renderCreateHook({ + initialProvider: mockProvider, + initialInstance: mockInstance, + }); + + act(() => result.current.handleNext()); // → step 2 + + act(() => { + result.current.setDeploymentName(" 1 Agent"); + result.current.setSelectedLlm("gpt-4"); + }); + + expect(result.current.canGoNext).toBe(false); + expect(result.current.isDeploymentNameValid).toBe(false); + expect(result.current.hasDeploymentNameFormatError).toBe(true); + }); + + it("allows when trimmed name starts with a unicode letter", () => { + const { result } = renderCreateHook({ + initialProvider: mockProvider, + initialInstance: mockInstance, + }); + + act(() => result.current.handleNext()); // → step 2 + + act(() => { + result.current.setDeploymentName(" Ágent"); + result.current.setSelectedLlm("gpt-4"); + }); + + expect(result.current.canGoNext).toBe(true); + expect(result.current.isDeploymentNameValid).toBe(true); + expect(result.current.hasDeploymentNameFormatError).toBe(false); + }); }); describe("Step 3 (Attach Flows)", () => { diff --git a/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/custom-tool-naming.test.tsx b/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/custom-tool-naming.test.tsx index 3acfacc87e..51328a4efe 100644 --- a/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/custom-tool-naming.test.tsx +++ b/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/custom-tool-naming.test.tsx @@ -180,4 +180,18 @@ describe("Custom tool naming", () => { "Tool Beta", ); }); + + it("rejects create payload when agent name does not start with a letter", () => { + const { result } = renderStepperHook(); + + act(() => { + result.current.setDeploymentName("1 Agent"); + result.current.setSelectedLlm("test-model"); + result.current.handleSelectVersion("flow-1", "ver-1", "v1"); + }); + + expect(() => result.current.buildDeploymentPayload("provider-1")).toThrow( + "Deployment name must start with a letter", + ); + }); }); diff --git a/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/edit-mode.test.tsx b/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/edit-mode.test.tsx index 86ca32849d..fd86cf3dfa 100644 --- a/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/edit-mode.test.tsx +++ b/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/edit-mode.test.tsx @@ -92,6 +92,31 @@ describe("Edit mode — basic state", () => { expect(result.current.canGoNext).toBe(true); }); + it("blocks update flow when existing name does not start with a letter", () => { + const invalidDeployment = { ...mockDeployment, name: "1 Agent" }; + const wrapper = ({ children }: { children: React.ReactNode }) => ( + + {children} + + ); + const { result } = renderHook(() => useDeploymentStepper(), { wrapper }); + + expect(result.current.canGoNext).toBe(false); + expect(result.current.isDeploymentNameValid).toBe(false); + expect(result.current.hasDeploymentNameFormatError).toBe(true); + expect(() => result.current.buildDeploymentUpdatePayload()).toThrow( + "Deployment name must start with a letter", + ); + }); + it("canGoNext on step 2 (Attach) allows proceeding in edit mode", () => { const { result } = renderEditHook(); act(() => result.current.handleNext()); // step 2 diff --git a/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/step-type.test.tsx b/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/step-type.test.tsx index aafd885939..1cd2f07022 100644 --- a/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/step-type.test.tsx +++ b/src/frontend/src/pages/MainPage/pages/deploymentsPage/__tests__/step-type.test.tsx @@ -13,6 +13,8 @@ const mockSetSelectedLlm = jest.fn(); let mockIsEditMode = false; let mockDeploymentType = "agent"; let mockDeploymentName = ""; +let mockIsDeploymentNameValid = false; +let mockHasDeploymentNameFormatError = false; let mockDeploymentDescription = ""; let mockSelectedLlm = ""; let mockSelectedInstance: { id: string } | null = { id: "inst-1" }; @@ -26,6 +28,8 @@ jest.mock("../contexts/deployment-stepper-context", () => ({ setDeploymentType: mockSetDeploymentType, deploymentName: mockDeploymentName, setDeploymentName: mockSetDeploymentName, + isDeploymentNameValid: mockIsDeploymentNameValid, + hasDeploymentNameFormatError: mockHasDeploymentNameFormatError, deploymentDescription: mockDeploymentDescription, setDeploymentDescription: mockSetDeploymentDescription, selectedLlm: mockSelectedLlm, @@ -61,6 +65,8 @@ beforeEach(() => { mockIsEditMode = false; mockDeploymentType = "agent"; mockDeploymentName = ""; + mockIsDeploymentNameValid = false; + mockHasDeploymentNameFormatError = false; mockDeploymentDescription = ""; mockSelectedLlm = ""; mockSelectedInstance = { id: "inst-1" }; @@ -143,6 +149,30 @@ describe("Name input", () => { screen.queryByText("Name cannot be changed after creation."), ).not.toBeInTheDocument(); }); + + it("shows validation error when name does not start with a letter", () => { + mockDeploymentName = "1 Agent"; + mockHasDeploymentNameFormatError = true; + render(); + expect( + screen.getByText("Agent name must start with a letter."), + ).toBeInTheDocument(); + expect(screen.getByPlaceholderText("e.g., Sales Bot")).toHaveAttribute( + "aria-invalid", + "true", + ); + }); + + it("does not show validation error for empty name", () => { + render(); + expect( + screen.queryByText("Agent name must start with a letter."), + ).not.toBeInTheDocument(); + expect(screen.getByPlaceholderText("e.g., Sales Bot")).toHaveAttribute( + "aria-invalid", + "false", + ); + }); }); // --------------------------------------------------------------------------- diff --git a/src/frontend/src/pages/MainPage/pages/deploymentsPage/components/step-type.tsx b/src/frontend/src/pages/MainPage/pages/deploymentsPage/components/step-type.tsx index a12764a652..22c4a38243 100644 --- a/src/frontend/src/pages/MainPage/pages/deploymentsPage/components/step-type.tsx +++ b/src/frontend/src/pages/MainPage/pages/deploymentsPage/components/step-type.tsx @@ -32,6 +32,7 @@ export default function StepType() { setDeploymentType, deploymentName, setDeploymentName, + hasDeploymentNameFormatError, deploymentDescription, setDeploymentDescription, selectedLlm, @@ -142,11 +143,21 @@ export default function StepType() { setDeploymentName(e.target.value)} disabled={isEditMode} + aria-invalid={hasDeploymentNameFormatError} /> + {hasDeploymentNameFormatError && ( + + Agent name must start with a letter. + + )} {isEditMode && ( Name cannot be changed after creation. diff --git a/src/frontend/src/pages/MainPage/pages/deploymentsPage/contexts/deployment-stepper-context.tsx b/src/frontend/src/pages/MainPage/pages/deploymentsPage/contexts/deployment-stepper-context.tsx index 79341c3be5..3a46d79cd4 100644 --- a/src/frontend/src/pages/MainPage/pages/deploymentsPage/contexts/deployment-stepper-context.tsx +++ b/src/frontend/src/pages/MainPage/pages/deploymentsPage/contexts/deployment-stepper-context.tsx @@ -73,6 +73,8 @@ interface DeploymentStepperContextType { setDeploymentType: (type: DeploymentType) => void; deploymentName: string; setDeploymentName: (name: string) => void; + isDeploymentNameValid: boolean; + hasDeploymentNameFormatError: boolean; deploymentDescription: string; setDeploymentDescription: (description: string) => void; selectedLlm: string; @@ -168,6 +170,11 @@ export function DeploymentStepperProvider({ >(initialState?.initialConnectionsByFlow ?? new Map()); const [hasToolNameErrors, setHasToolNameErrors] = useState(false); + const trimmedDeploymentName = deploymentName.trim(); + const hasDeploymentNameFormatError = + trimmedDeploymentName !== "" && !/^\p{L}/u.test(trimmedDeploymentName); + const isDeploymentNameValid = + trimmedDeploymentName !== "" && !hasDeploymentNameFormatError; // Edit mode: track which pre-existing flows the user wants to detach. const [removedFlowIds, setRemovedFlowIds] = useState>(new Set()); @@ -253,7 +260,7 @@ export function DeploymentStepperProvider({ ); } if (logical === 2) { - return deploymentName.trim() !== "" && selectedLlm.trim() !== ""; + return isDeploymentNameValid && selectedLlm.trim() !== ""; } if (logical === 3) { // In edit mode, user can proceed without new attachments (may just change desc/LLM). @@ -270,6 +277,7 @@ export function DeploymentStepperProvider({ selectedInstance, hasValidCredentials, deploymentName, + isDeploymentNameValid, selectedLlm, selectedVersionByFlow, isEditMode, @@ -357,6 +365,9 @@ export function DeploymentStepperProvider({ const buildDeploymentPayload = useCallback( (providerId: string): DeploymentCreateRequest => { + if (!isDeploymentNameValid) { + throw new Error("Deployment name must start with a letter"); + } const allConnectionIds = new Set(); Array.from(attachedConnectionByFlow.values()).forEach((ids) => { ids.forEach((id) => allConnectionIds.add(id)); @@ -381,7 +392,7 @@ export function DeploymentStepperProvider({ ...(initialState?.projectId ? { project_id: initialState.projectId } : {}), - name: deploymentName, + name: trimmedDeploymentName, description: deploymentDescription, type: deploymentType, provider_data: { @@ -396,10 +407,11 @@ export function DeploymentStepperProvider({ buildConnectionPayloads, initialState?.projectId, deploymentDescription, - deploymentName, deploymentType, + isDeploymentNameValid, selectedLlm, selectedVersionByFlow, + trimmedDeploymentName, toolNameByFlow, ], ); @@ -411,6 +423,9 @@ export function DeploymentStepperProvider({ "buildDeploymentUpdatePayload called outside edit mode", ); } + if (!isDeploymentNameValid) { + throw new Error("Deployment name must start with a letter"); + } const result: DeploymentUpdateRequest = { deployment_id: editingDeployment.id, @@ -510,6 +525,7 @@ export function DeploymentStepperProvider({ }, [ editingDeployment, deploymentDescription, + isDeploymentNameValid, selectedLlm, initialVersionByFlow, initialToolNameByFlow, @@ -541,6 +557,8 @@ export function DeploymentStepperProvider({ setDeploymentType, deploymentName, setDeploymentName, + isDeploymentNameValid, + hasDeploymentNameFormatError, deploymentDescription, setDeploymentDescription, selectedLlm, @@ -581,6 +599,8 @@ export function DeploymentStepperProvider({ credentials, deploymentType, deploymentName, + isDeploymentNameValid, + hasDeploymentNameFormatError, deploymentDescription, selectedLlm, connections, From 7a84f8b254519816a2b3e271ac47972299ff0f95 Mon Sep 17 00:00:00 2001 From: Eric Hare Date: Mon, 27 Apr 2026 17:29:58 -0700 Subject: [PATCH 2/3] fix(mcp): enforce auth on /streamable for auth_type=oauth projects (#12756) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(mcp): enforce authentication on /streamable for auth_type=oauth projects Previously, projects configured with auth_type=oauth did not enforce authentication on the /streamable (and /sse) endpoints when MCP Composer was enabled. verify_project_auth only enforced API key authentication for auth_type=apikey; for oauth it fell through to the superuser fallback, returning HTTP 200 with full initialize/tools/list payloads to unauthenticated requests. This treats auth_type=oauth the same as auth_type=apikey on the direct langflow MCP transport endpoints: requests must present a valid API key or they are rejected with HTTP 401. End users of OAuth-configured projects should continue to connect via the configured MCP Composer OAuth endpoint (which proxies to these endpoints with backend credentials). * fix(mcp): preserve composer loopback forwarding on OAuth auth enforcement Address review feedback on PR #12756. Requiring an API key for every OAuth request would have broken the intended MCP Composer proxy path: the composer subprocess is started with only the Langflow endpoint URL and OAuth env vars and does not currently inject a backend x-api-key when it forwards authenticated OAuth traffic, so hardening /streamable and /sse would turn every legitimate composer-forwarded request into a 401. Keep the enforcement against direct unauthenticated external access, but trust requests whose direct TCP peer is a loopback address (the on-host MCP Composer subprocess). Remote callers still need a valid x-api-key to hit the project transport endpoints when auth_type=oauth. Add coverage for the loopback passthrough and the _is_loopback_client helper, and gate the rejection tests on a patched non-loopback client so they exercise the remote path even though httpx AsyncClient reports 127.0.0.1 by default. * fix(mcp): drop loopback trust shortcut; always require API key for OAuth transport Follow-up to PR #12756. The loopback-based passthrough from the prior commit is unsafe in deployments where Langflow sits behind a same-host reverse proxy or sidecar: the proxy becomes the direct TCP peer, so external unauthenticated traffic satisfies the loopback check and falls through to the superuser, reopening the original bypass. There is no composer-specific identity on the wire today — mcp-composer forwards without any Langflow credential — so loopback address cannot distinguish the trusted subprocess from another loopback peer. Require a valid x-api-key on /streamable and /sse for every OAuth project regardless of source. Until mcp-composer can forward a project-scoped backend credential, the composer proxy path will return 401; this is the intended secure behavior. Remove the _is_loopback_client helper, the client_host plumbing into verify_project_auth, and the associated tests. Drop the force_non_loopback_client fixture; the rejection tests now exercise the real code path directly. * fix(mcp): clarify OAuth 401 detail about the current API-key requirement The previous message asked callers to connect through MCP Composer, but the composer proxy path currently returns 401 itself (see PR #12756 discussion) because mcp-composer cannot forward a backend credential yet. Describe the actual working path — use an x-api-key — and note that composer-based access is pending credential forwarding support. * fix: No timeout when connect to oauth project * Secrets baseline update * fix: Standardize paths on Windows --- .secrets.baseline | 10 +- .../base/langflow/api/v1/mcp_projects.py | 66 ++- .../tests/unit/api/v1/test_mcp_projects.py | 143 +++++++ .../src/lfx/services/mcp_composer/service.py | 375 +++++++++++------- .../services/settings/test_mcp_composer.py | 76 +++- 5 files changed, 519 insertions(+), 151 deletions(-) diff --git a/.secrets.baseline b/.secrets.baseline index de765fa22f..23582936a2 100644 --- a/.secrets.baseline +++ b/.secrets.baseline @@ -2793,7 +2793,7 @@ "filename": "src/backend/tests/unit/api/v1/test_mcp_projects.py", "hashed_secret": "4258d43e3b1f9658067ceea9c682a96cbdbb5ca0", "is_verified": false, - "line_number": 596, + "line_number": 739, "is_secret": false } ], @@ -8101,7 +8101,7 @@ "filename": "src/lfx/tests/unit/services/settings/test_mcp_composer.py", "hashed_secret": "00942f4668670f34c5943cf52c7ef3139fe2b8d6", "is_verified": false, - "line_number": 162, + "line_number": 236, "is_secret": false }, { @@ -8109,7 +8109,7 @@ "filename": "src/lfx/tests/unit/services/settings/test_mcp_composer.py", "hashed_secret": "a94a8fe5ccb19ba61c4c0873d391e987982fbbd3", "is_verified": false, - "line_number": 202, + "line_number": 276, "is_secret": false }, { @@ -8117,7 +8117,7 @@ "filename": "src/lfx/tests/unit/services/settings/test_mcp_composer.py", "hashed_secret": "e5e9fa1ba31ecd1ae84f75caaa474f3a663f05f4", "is_verified": false, - "line_number": 335, + "line_number": 409, "is_secret": false } ], @@ -8198,5 +8198,5 @@ } ] }, - "generated_at": "2026-04-24T03:31:30Z" + "generated_at": "2026-04-27T15:53:27Z" } diff --git a/src/backend/base/langflow/api/v1/mcp_projects.py b/src/backend/base/langflow/api/v1/mcp_projects.py index 71cb7b77d9..3d1b44049f 100644 --- a/src/backend/base/langflow/api/v1/mcp_projects.py +++ b/src/backend/base/langflow/api/v1/mcp_projects.py @@ -21,7 +21,11 @@ from lfx.base.mcp.constants import MAX_MCP_SERVER_NAME_LENGTH from lfx.base.mcp.util import sanitize_mcp_name from lfx.log import logger from lfx.services.deps import get_settings_service, session_scope -from lfx.services.mcp_composer.service import MCPComposerError, MCPComposerService +from lfx.services.mcp_composer.service import ( + COMPOSER_BACKEND_AUTH_HEADER, + MCPComposerError, + MCPComposerService, +) from lfx.services.schema import ServiceType from mcp import types from mcp.server import NotificationOptions, Server @@ -66,7 +70,7 @@ from langflow.services.auth.constants import AUTO_LOGIN_WARNING from langflow.services.auth.mcp_encryption import decrypt_auth_settings, encrypt_auth_settings from langflow.services.database.models import Flow, Folder from langflow.services.database.models.api_key.crud import check_key, create_api_key -from langflow.services.database.models.api_key.model import ApiKey, ApiKeyCreate +from langflow.services.database.models.api_key.model import ApiKeyCreate from langflow.services.database.models.user.crud import get_user_by_username from langflow.services.database.models.user.model import User from langflow.services.deps import get_service @@ -80,8 +84,9 @@ router = APIRouter(prefix="/mcp/project", tags=["mcp_projects"]) async def verify_project_auth( db: AsyncSession, project_id: UUID, - query_param: str, - header_param: str, + query_param: str | None, + header_param: str | None, + composer_backend_token: str | None = None, ) -> User: """MCP-specific user authentication that allows fallback to username lookup when not using API key auth. @@ -89,7 +94,6 @@ async def verify_project_auth( or checks if the API key is valid. """ settings_service = get_settings_service() - result: ApiKey | User | None project = (await db.exec(select(Folder).where(Folder.id == project_id))).first() @@ -101,14 +105,42 @@ async def verify_project_auth( if project.auth_settings: auth_settings = AuthSettings(**project.auth_settings) - if (not auth_settings and not settings_service.auth_settings.AUTO_LOGIN) or ( - auth_settings and auth_settings.auth_type == "apikey" - ): + project_auth_type = auth_settings.auth_type if auth_settings else None + if project_auth_type == "oauth" and composer_backend_token: + mcp_composer_service: MCPComposerService = cast( + MCPComposerService, get_service(ServiceType.MCP_COMPOSER_SERVICE) + ) + if mcp_composer_service.validate_backend_auth_token(str(project_id), composer_backend_token): + if project.user_id: + project_user = await db.get(User, project.user_id) + if project_user: + return project_user + raise HTTPException(status_code=404, detail="Project owner not found") + + # OAuth projects must present a valid API key at the Langflow transport endpoint: network-level + # trust (loopback / same-host proxy) is unsafe because it cannot distinguish the local MCP + # Composer subprocess from another loopback peer behind a reverse proxy or sidecar. The + # composer-to-Langflow hop should be authenticated explicitly once mcp-composer can forward + # a project-scoped backend credential; until then, direct backend access requires a key. + requires_api_key = (not auth_settings and not settings_service.auth_settings.AUTO_LOGIN) or ( + project_auth_type in {"apikey", "oauth"} + ) + + if requires_api_key: api_key = query_param or header_param if not api_key: + if project_auth_type == "oauth": + detail = ( + "This project is configured for OAuth authentication, but the MCP transport endpoint " + "currently requires a valid x-api-key header or query parameter for backend access. " + "Credential forwarding from MCP Composer is not yet available; use an API key in the " + "meantime." + ) + else: + detail = "API key required for this project. Provide x-api-key header or query parameter." raise HTTPException( status_code=401, - detail="API key required for this project. Provide x-api-key header or query parameter.", + detail=detail, ) # Validate the API key @@ -126,13 +158,16 @@ async def verify_project_auth( return user - # Get the first user + return await _superuser_fallback(db, settings_service) + + +async def _superuser_fallback(db: AsyncSession, settings_service) -> User: + """Resolve the configured superuser for unauthenticated MCP paths that allow fallback.""" if not settings_service.auth_settings.SUPERUSER: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Missing superuser username in auth settings", ) - # For MCP endpoints, always fall back to username lookup when no API key is provided result = await get_user_by_username(db, settings_service.auth_settings.SUPERUSER) if result: logger.warning(AUTO_LOGIN_WARNING) @@ -169,10 +204,17 @@ async def verify_project_auth_conditional( # Extract API keys api_key_query_value = request.query_params.get("x-api-key") api_key_header_value = request.headers.get("x-api-key") + composer_backend_token = request.headers.get(COMPOSER_BACKEND_AUTH_HEADER) # Check if this project requires API key only authentication if get_settings_service().settings.mcp_composer_enabled: - return await verify_project_auth(session, project_id, api_key_query_value, api_key_header_value) + return await verify_project_auth( + session, + project_id, + api_key_query_value, + api_key_header_value, + composer_backend_token, + ) # For all other cases, use standard MCP authentication (allows JWT + API keys) # Call the MCP auth function directly diff --git a/src/backend/tests/unit/api/v1/test_mcp_projects.py b/src/backend/tests/unit/api/v1/test_mcp_projects.py index d7a6040d36..2935d4b14c 100644 --- a/src/backend/tests/unit/api/v1/test_mcp_projects.py +++ b/src/backend/tests/unit/api/v1/test_mcp_projects.py @@ -26,6 +26,7 @@ from langflow.services.deps import get_settings_service from lfx.base.mcp.constants import MAX_MCP_SERVER_NAME_LENGTH from lfx.base.mcp.util import sanitize_mcp_name from lfx.services.deps import session_scope +from lfx.services.mcp_composer.service import COMPOSER_BACKEND_AUTH_HEADER from mcp.server.sse import SseServerTransport from sqlmodel import select @@ -262,6 +263,148 @@ async def test_handle_project_streamable_messages_success( mock_streamable_http_manager.handle_request.assert_called_once() +async def _set_project_auth_type(project_id, auth_type: str) -> None: + """Persist an auth_settings value for the given project.""" + from langflow.services.auth.mcp_encryption import encrypt_auth_settings + + async with session_scope() as session: + project = await session.get(Folder, project_id) + assert project is not None + project.auth_settings = encrypt_auth_settings({"auth_type": auth_type}) + session.add(project) + + +async def test_streamable_rejects_unauthenticated_oauth_project( + client: AsyncClient, + user_test_project, + mock_streamable_http_manager, + enable_mcp_composer, +): + """OAuth projects must reject any unauthenticated /streamable request. + + Network-level trust (loopback / same-host proxy) is intentionally NOT used here: a + same-host reverse proxy or sidecar would make every external request appear to be + loopback, which would reopen the original unauthenticated bypass. Requests must + present a valid x-api-key regardless of source. + """ + assert enable_mcp_composer + await _set_project_auth_type(user_test_project.id, "oauth") + + response = await client.post( + f"api/v1/mcp/project/{user_test_project.id}/streamable", + json={"type": "test", "content": "message"}, + ) + + assert response.status_code == status.HTTP_401_UNAUTHORIZED + assert "OAuth" in response.json()["detail"] + mock_streamable_http_manager.handle_request.assert_not_called() + + +async def test_streamable_rejects_unauthenticated_oauth_project_trailing_slash( + client: AsyncClient, + user_test_project, + mock_streamable_http_manager, + enable_mcp_composer, +): + """Trailing-slash variant of /streamable must also enforce OAuth auth.""" + assert enable_mcp_composer + await _set_project_auth_type(user_test_project.id, "oauth") + + response = await client.post( + f"api/v1/mcp/project/{user_test_project.id}/streamable/", + json={"type": "test", "content": "message"}, + ) + + assert response.status_code == status.HTTP_401_UNAUTHORIZED + mock_streamable_http_manager.handle_request.assert_not_called() + + +async def test_sse_rejects_unauthenticated_oauth_project( + client: AsyncClient, + user_test_project, + mock_sse_transport, + enable_mcp_composer, +): + """SSE endpoint must also reject unauthenticated OAuth-project requests.""" + assert enable_mcp_composer + await _set_project_auth_type(user_test_project.id, "oauth") + + response = await client.get(f"api/v1/mcp/project/{user_test_project.id}/sse") + + assert response.status_code == status.HTTP_401_UNAUTHORIZED + mock_sse_transport.connect_sse.assert_not_called() + + +async def test_streamable_oauth_project_accepts_valid_api_key( + client: AsyncClient, + user_test_project, + created_api_key, + mock_streamable_http_manager, + enable_mcp_composer, +): + """Valid API keys must be accepted for OAuth-configured projects.""" + assert enable_mcp_composer + await _set_project_auth_type(user_test_project.id, "oauth") + + response = await client.post( + f"api/v1/mcp/project/{user_test_project.id}/streamable", + headers={"x-api-key": created_api_key.api_key}, + json={"type": "test", "content": "message"}, + ) + + assert response.status_code == status.HTTP_200_OK + mock_streamable_http_manager.handle_request.assert_called_once() + + +async def test_streamable_oauth_project_accepts_valid_composer_backend_token( + client: AsyncClient, + user_test_project, + mock_streamable_http_manager, + enable_mcp_composer, +): + """The in-process MCP Composer token must authenticate Composer's backend hop.""" + assert enable_mcp_composer + await _set_project_auth_type(user_test_project.id, "oauth") + + composer_service = MagicMock() + composer_service.validate_backend_auth_token.return_value = True + + with patch("langflow.api.v1.mcp_projects.get_service", return_value=composer_service): + response = await client.post( + f"api/v1/mcp/project/{user_test_project.id}/streamable", + headers={COMPOSER_BACKEND_AUTH_HEADER: "valid-composer-token"}, + json={"type": "test", "content": "message"}, + ) + + assert response.status_code == status.HTTP_200_OK + composer_service.validate_backend_auth_token.assert_called_once_with( + str(user_test_project.id), + "valid-composer-token", + ) + mock_streamable_http_manager.handle_request.assert_called_once() + + +async def test_streamable_oauth_project_rejects_invalid_api_key( + client: AsyncClient, + user_test_project, + mock_streamable_http_manager, + enable_mcp_composer, +): + """Invalid API keys must be rejected for OAuth-configured projects.""" + assert enable_mcp_composer + await _set_project_auth_type(user_test_project.id, "oauth") + + response = await client.post( + f"api/v1/mcp/project/{user_test_project.id}/streamable", + headers={"x-api-key": "not-a-real-api-key"}, + json={"type": "test", "content": "message"}, + ) + + assert response.status_code == status.HTTP_401_UNAUTHORIZED + assert response.json()["detail"] == "Invalid API key" + mock_streamable_http_manager.handle_request.assert_not_called() + + async def test_handle_project_messages_success( client: AsyncClient, user_test_project, mock_sse_transport, logged_in_headers ): diff --git a/src/lfx/src/lfx/services/mcp_composer/service.py b/src/lfx/src/lfx/services/mcp_composer/service.py index 2727bcd3fd..da806b11be 100644 --- a/src/lfx/src/lfx/services/mcp_composer/service.py +++ b/src/lfx/src/lfx/services/mcp_composer/service.py @@ -5,6 +5,7 @@ import json import os import platform import re +import secrets import select import socket import subprocess @@ -22,6 +23,7 @@ from lfx.services.deps import get_settings_service GENERIC_STARTUP_ERROR_MSG = ( "MCP Composer startup failed. Check OAuth configuration and check logs for more information." ) +COMPOSER_BACKEND_AUTH_HEADER = "x-langflow-mcp-composer-token" class MCPComposerError(Exception): @@ -85,6 +87,7 @@ class MCPComposerService(Service): self._port_to_project: dict[int, str] = {} # Track which project is using which port self._pid_to_project: dict[int, str] = {} # Track which PID belongs to which project self._last_errors: dict[str, str] = {} # Track last error message per project for UI display + self._backend_auth_tokens: dict[str, str] = {} # project_id -> internal Composer-to-Langflow token def get_last_error(self, project_id: str) -> str | None: """Get the last error message for a project, if any.""" @@ -98,6 +101,19 @@ class MCPComposerService(Service): """Clear the last error message for a project.""" self._last_errors.pop(project_id, None) + def get_or_create_backend_auth_token(self, project_id: str) -> str: + """Return the internal token used by MCP Composer to call Langflow's MCP endpoint.""" + if project_id not in self._backend_auth_tokens: + self._backend_auth_tokens[project_id] = secrets.token_urlsafe(32) + return self._backend_auth_tokens[project_id] + + def validate_backend_auth_token(self, project_id: str, token: str | None) -> bool: + """Validate an internal Composer-to-Langflow token for a project.""" + expected_token = self._backend_auth_tokens.get(project_id) + if not expected_token or not token: + return False + return secrets.compare_digest(expected_token, token) + def _is_port_available(self, port: int, host: str = "localhost") -> bool: """Check if a port is available by trying to bind to it. @@ -531,6 +547,7 @@ class MCPComposerService(Service): # Remove from tracking self.project_composers.pop(project_id, None) + self._backend_auth_tokens.pop(project_id, None) await logger.adebug(f"Removed tracking for project {project_id}") async def _wait_for_process_exit(self, process): @@ -563,6 +580,8 @@ class MCPComposerService(Service): try: # On Windows with temp files, read from files instead of pipes if stdout_file and stderr_file: + stdout_path = Path(stdout_file.name) + stderr_path = Path(stderr_file.name) # Close file handles to flush and allow reading try: stdout_file.close() @@ -576,7 +595,7 @@ class MCPComposerService(Service): def read_file(filepath): return Path(filepath).read_bytes() - stdout_bytes = await asyncio.to_thread(read_file, stdout_file.name) + stdout_bytes = await asyncio.to_thread(read_file, stdout_path) stdout_content = stdout_bytes.decode("utf-8", errors="replace") if stdout_bytes else "" except Exception as e: # noqa: BLE001 await logger.adebug(f"Error reading stdout file: {e}") @@ -586,15 +605,15 @@ class MCPComposerService(Service): def read_file(filepath): return Path(filepath).read_bytes() - stderr_bytes = await asyncio.to_thread(read_file, stderr_file.name) + stderr_bytes = await asyncio.to_thread(read_file, stderr_path) stderr_content = stderr_bytes.decode("utf-8", errors="replace") if stderr_bytes else "" except Exception as e: # noqa: BLE001 await logger.adebug(f"Error reading stderr file: {e}") # Clean up temp files try: - Path(stdout_file.name).unlink() - Path(stderr_file.name).unlink() + stdout_path.unlink() + stderr_path.unlink() except Exception as e: # noqa: BLE001 await logger.adebug(f"Error removing temp files: {e}") else: @@ -822,6 +841,67 @@ class MCPComposerService(Service): config_error_msg = f"Invalid OAuth configuration: {'; '.join(error_parts)}" raise MCPComposerConfigError(config_error_msg) + def _build_backend_auth_member_config( + self, + project_id: str, + streamable_http_url: str, + *, + legacy_sse_url: str | None = None, + ) -> list[dict[str, Any]]: + """Build mcp-composer member-server config with internal Langflow auth headers.""" + backend_token = self.get_or_create_backend_auth_token(project_id) + headers = {COMPOSER_BACKEND_AUTH_HEADER: backend_token} + streamable_id = f"{project_id}-streamable" + config: list[dict[str, Any]] = [ + { + "id": streamable_id, + "_id": streamable_id, + "type": "http", + "endpoint": streamable_http_url, + "headers": headers, + } + ] + + if legacy_sse_url: + sse_id = f"{project_id}-sse" + config.append( + { + "id": sse_id, + "_id": sse_id, + "type": "sse", + "endpoint": legacy_sse_url, + "headers": headers, + } + ) + + return config + + def _write_backend_auth_member_config( + self, + project_id: str, + streamable_http_url: str, + *, + legacy_sse_url: str | None = None, + ) -> Path: + """Write a temporary mcp-composer config containing backend auth headers.""" + config = self._build_backend_auth_member_config( + project_id, + streamable_http_url, + legacy_sse_url=legacy_sse_url, + ) + fd, config_path = tempfile.mkstemp( + prefix=f"mcp_composer_{project_id}_config_", + suffix=".json", + ) + path = Path(config_path) + try: + with os.fdopen(fd, mode="w", encoding="utf-8") as config_file: + json.dump(config, config_file) + except Exception: + path.unlink(missing_ok=True) + raise + return path + @staticmethod def _normalize_config_value(value: Any) -> Any: """Normalize a config value (None or empty string becomes None). @@ -1280,6 +1360,8 @@ class MCPComposerService(Service): settings = get_settings_service().settings # Some composer tooling still uses the --sse-url flag for backwards compatibility even in HTTP mode. effective_legacy_sse_url = legacy_sse_url or f"{streamable_http_url.rstrip('/')}/sse" + auth_type = auth_config.get("auth_type") if auth_config else None + backend_auth_config_path: Path | None = None cmd = [ "uvx", @@ -1290,53 +1372,64 @@ class MCPComposerService(Service): host, "--mode", "http", - "--endpoint", - streamable_http_url, - "--sse-url", - effective_legacy_sse_url, - "--disable-composer-tools", ] + if auth_type == "oauth": + backend_auth_config_path = self._write_backend_auth_member_config( + project_id, + streamable_http_url, + legacy_sse_url=effective_legacy_sse_url, + ) + cmd.extend(["--config_path", str(backend_auth_config_path)]) + else: + cmd.extend( + [ + "--endpoint", + streamable_http_url, + "--sse-url", + effective_legacy_sse_url, + ] + ) + + cmd.append("--disable-composer-tools") + # Set environment variables env = os.environ.copy() oauth_server_url = auth_config.get("oauth_server_url") if auth_config else None - if auth_config: - auth_type = auth_config.get("auth_type") + if auth_config and auth_type == "oauth": + cmd.extend(["--auth_type", "oauth"]) - if auth_type == "oauth": - cmd.extend(["--auth_type", "oauth"]) + # Add OAuth environment variables as command line arguments + cmd.extend(["--env", "ENABLE_OAUTH", "True"]) - # Add OAuth environment variables as command line arguments - cmd.extend(["--env", "ENABLE_OAUTH", "True"]) + # Map auth config to environment variables for OAuth + # Note: oauth_host and oauth_port are passed both via --host/--port CLI args + # (for server binding) and as environment variables (for OAuth flow) + # Note: mcp-composer expects the env var name OAUTH_CALLBACK_PATH, but the value + # is still the full callback URL used as the OAuth redirect_uri. Support legacy + # oauth_callback_path input for backwards compatibility. + oauth_env_mapping = { + "oauth_host": "OAUTH_HOST", + "oauth_port": "OAUTH_PORT", + "oauth_server_url": "OAUTH_SERVER_URL", + "oauth_callback_path": "OAUTH_CALLBACK_PATH", + "oauth_client_id": "OAUTH_CLIENT_ID", + "oauth_client_secret": "OAUTH_CLIENT_SECRET", # pragma: allowlist secret + "oauth_auth_url": "OAUTH_AUTH_URL", + "oauth_token_url": "OAUTH_TOKEN_URL", + "oauth_mcp_scope": "OAUTH_MCP_SCOPE", + "oauth_provider_scope": "OAUTH_PROVIDER_SCOPE", + } - # Map auth config to environment variables for OAuth - # Note: oauth_host and oauth_port are passed both via --host/--port CLI args - # (for server binding) and as environment variables (for OAuth flow) - # Note: mcp-composer expects the env var name OAUTH_CALLBACK_PATH, but the value - # is still the full callback URL used as the OAuth redirect_uri. Support legacy - # oauth_callback_path input for backwards compatibility. - oauth_env_mapping = { - "oauth_host": "OAUTH_HOST", - "oauth_port": "OAUTH_PORT", - "oauth_server_url": "OAUTH_SERVER_URL", - "oauth_callback_path": "OAUTH_CALLBACK_PATH", - "oauth_client_id": "OAUTH_CLIENT_ID", - "oauth_client_secret": "OAUTH_CLIENT_SECRET", # pragma: allowlist secret - "oauth_auth_url": "OAUTH_AUTH_URL", - "oauth_token_url": "OAUTH_TOKEN_URL", - "oauth_mcp_scope": "OAUTH_MCP_SCOPE", - "oauth_provider_scope": "OAUTH_PROVIDER_SCOPE", - } + normalized_auth_config = self._normalize_oauth_callback_aliases(auth_config) - normalized_auth_config = self._normalize_oauth_callback_aliases(auth_config) - - # Add environment variables as command line arguments - # Only set non-empty values to avoid Pydantic validation errors - for config_key, env_key in oauth_env_mapping.items(): - value = normalized_auth_config.get(config_key) - if value is not None and str(value).strip(): - cmd.extend(["--env", env_key, str(value)]) + # Add environment variables as command line arguments + # Only set non-empty values to avoid Pydantic validation errors + for config_key, env_key in oauth_env_mapping.items(): + value = normalized_auth_config.get(config_key) + if value is not None and str(value).strip(): + cmd.extend(["--env", env_key, str(value)]) # Log the command being executed (with secrets obfuscated) safe_cmd = self._obfuscate_command_secrets(cmd) @@ -1348,6 +1441,8 @@ class MCPComposerService(Service): stderr_handle: int | typing.IO[bytes] = subprocess.PIPE stdout_file = None stderr_file = None + stdout_path: Path | None = None + stderr_path: Path | None = None if platform.system() == "Windows": # Create temp files for stdout/stderr on Windows to avoid pipe deadlocks @@ -1361,30 +1456,89 @@ class MCPComposerService(Service): ) stdout_handle = stdout_file stderr_handle = stderr_file - stdout_name = stdout_file.name - stderr_name = stderr_file.name - await logger.adebug(f"Using temp files for MCP Composer logs: stdout={stdout_name}, stderr={stderr_name}") - - process = subprocess.Popen(cmd, env=env, stdout=stdout_handle, stderr=stderr_handle) # noqa: ASYNC220, S603 - - # Monitor the process startup with multiple checks - process_running = False - port_bound = False - - await logger.adebug( - f"MCP Composer process started with PID {process.pid}, monitoring startup for project {project_id}..." - ) + stdout_path = Path(stdout_file.name) + stderr_path = Path(stderr_file.name) + await logger.adebug(f"Using temp files for MCP Composer logs: stdout={stdout_path}, stderr={stderr_path}") try: - for check in range(max_startup_checks): - await asyncio.sleep(startup_delay) + process = subprocess.Popen(cmd, env=env, stdout=stdout_handle, stderr=stderr_handle) # noqa: ASYNC220, S603 - # Check if process is still running + # Monitor the process startup with multiple checks + process_running = False + port_bound = False + + await logger.adebug( + f"MCP Composer process started with PID {process.pid}, monitoring startup for project {project_id}..." + ) + + try: + for check in range(max_startup_checks): + await asyncio.sleep(startup_delay) + + # Check if process is still running + poll_result = process.poll() + + startup_error_msg = None + if poll_result is not None: + # Process terminated, get the error output + ( + stdout_content, + stderr_content, + startup_error_msg, + ) = await self._read_process_output_and_extract_error( + process, oauth_server_url, stdout_file=stdout_file, stderr_file=stderr_file + ) + await self._log_startup_error_details( + project_id, cmd, host, port, stdout_content, stderr_content, startup_error_msg, poll_result + ) + raise MCPComposerStartupError(startup_error_msg, project_id) + + # Process is still running, check if port is bound + port_bound = not self._is_port_available(port) + + if port_bound: + await logger.adebug( + f"MCP Composer for project {project_id} bound to port {port} " + f"(check {check + 1}/{max_startup_checks})" + ) + process_running = True + break + await logger.adebug( + f"MCP Composer for project {project_id} not yet bound to port {port} " + f"(check {check + 1}/{max_startup_checks})" + ) + + # Try to read any available stderr/stdout without blocking to see what's happening + await self._read_stream_non_blocking(process.stderr, "stderr") + await self._read_stream_non_blocking(process.stdout, "stdout") + + except asyncio.CancelledError: + # Operation was cancelled, kill the process and cleanup + await logger.adebug( + f"MCP Composer process startup cancelled for project {project_id}, " + f"terminating process {process.pid}" + ) + try: + process.terminate() + # Wait for graceful termination with timeout + try: + await asyncio.wait_for(asyncio.to_thread(process.wait), timeout=2.0) + except asyncio.TimeoutError: + # Force kill if graceful termination times out + await logger.adebug(f"Process {process.pid} did not terminate gracefully, force killing") + await asyncio.to_thread(process.kill) + await asyncio.to_thread(process.wait) + except Exception as e: # noqa: BLE001 + await logger.adebug(f"Error terminating process during cancellation: {e}") + raise # Re-raise to propagate cancellation + + # After all checks + if not process_running or not port_bound: + # Get comprehensive error information poll_result = process.poll() - startup_error_msg = None if poll_result is not None: - # Process terminated, get the error output + # Process died ( stdout_content, stderr_content, @@ -1396,91 +1550,46 @@ class MCPComposerService(Service): project_id, cmd, host, port, stdout_content, stderr_content, startup_error_msg, poll_result ) raise MCPComposerStartupError(startup_error_msg, project_id) - - # Process is still running, check if port is bound - port_bound = not self._is_port_available(port) - - if port_bound: - await logger.adebug( - f"MCP Composer for project {project_id} bound to port {port} " - f"(check {check + 1}/{max_startup_checks})" - ) - process_running = True - break - await logger.adebug( - f"MCP Composer for project {project_id} not yet bound to port {port} " - f"(check {check + 1}/{max_startup_checks})" + # Process running but port not bound + await logger.aerror( + f" - Checked {max_startup_checks} times over {max_startup_checks * startup_delay} seconds" ) - # Try to read any available stderr/stdout without blocking to see what's happening - await self._read_stream_non_blocking(process.stderr, "stderr") - await self._read_stream_non_blocking(process.stdout, "stdout") - - except asyncio.CancelledError: - # Operation was cancelled, kill the process and cleanup - await logger.adebug( - f"MCP Composer process startup cancelled for project {project_id}, terminating process {process.pid}" - ) - try: + # Get any available output before terminating process.terminate() - # Wait for graceful termination with timeout - try: - await asyncio.wait_for(asyncio.to_thread(process.wait), timeout=2.0) - except asyncio.TimeoutError: - # Force kill if graceful termination times out - await logger.adebug(f"Process {process.pid} did not terminate gracefully, force killing") - await asyncio.to_thread(process.kill) - await asyncio.to_thread(process.wait) - except Exception as e: # noqa: BLE001 - await logger.adebug(f"Error terminating process during cancellation: {e}") - raise # Re-raise to propagate cancellation - - # After all checks - if not process_running or not port_bound: - # Get comprehensive error information - poll_result = process.poll() - - if poll_result is not None: - # Process died stdout_content, stderr_content, startup_error_msg = await self._read_process_output_and_extract_error( process, oauth_server_url, stdout_file=stdout_file, stderr_file=stderr_file ) await self._log_startup_error_details( - project_id, cmd, host, port, stdout_content, stderr_content, startup_error_msg, poll_result + project_id, cmd, host, port, stdout_content, stderr_content, startup_error_msg, pid=process.pid ) raise MCPComposerStartupError(startup_error_msg, project_id) - # Process running but port not bound - await logger.aerror( - f" - Checked {max_startup_checks} times over {max_startup_checks * startup_delay} seconds" - ) - # Get any available output before terminating - process.terminate() - stdout_content, stderr_content, startup_error_msg = await self._read_process_output_and_extract_error( - process, oauth_server_url, stdout_file=stdout_file, stderr_file=stderr_file - ) - await self._log_startup_error_details( - project_id, cmd, host, port, stdout_content, stderr_content, startup_error_msg, pid=process.pid - ) - raise MCPComposerStartupError(startup_error_msg, project_id) + # Close the pipes/files if everything is successful + if stdout_file and stderr_file: + # Clean up temp files on success + try: + stdout_file.close() + stderr_file.close() + if stdout_path: + stdout_path.unlink(missing_ok=True) + if stderr_path: + stderr_path.unlink(missing_ok=True) + except Exception as e: # noqa: BLE001 + await logger.adebug(f"Error cleaning up temp files on success: {e}") + else: + if process.stdout: + process.stdout.close() + if process.stderr: + process.stderr.close() - # Close the pipes/files if everything is successful - if stdout_file and stderr_file: - # Clean up temp files on success - try: - stdout_file.close() - stderr_file.close() - Path(stdout_file.name).unlink() - Path(stderr_file.name).unlink() - except Exception as e: # noqa: BLE001 - await logger.adebug(f"Error cleaning up temp files on success: {e}") - else: - if process.stdout: - process.stdout.close() - if process.stderr: - process.stderr.close() - - return process + return process + finally: + if backend_auth_config_path: + try: + backend_auth_config_path.unlink(missing_ok=True) + except Exception as e: # noqa: BLE001 + await logger.adebug(f"Error cleaning up MCP Composer config file: {e}") @require_composer_enabled def get_project_composer_port(self, project_id: str) -> int | None: diff --git a/src/lfx/tests/unit/services/settings/test_mcp_composer.py b/src/lfx/tests/unit/services/settings/test_mcp_composer.py index 5585accf81..089e016077 100644 --- a/src/lfx/tests/unit/services/settings/test_mcp_composer.py +++ b/src/lfx/tests/unit/services/settings/test_mcp_composer.py @@ -3,10 +3,15 @@ import asyncio import contextlib import socket +from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import pytest -from lfx.services.mcp_composer.service import MCPComposerPortError, MCPComposerService +from lfx.services.mcp_composer.service import ( + COMPOSER_BACKEND_AUTH_HEADER, + MCPComposerPortError, + MCPComposerService, +) @pytest.fixture @@ -38,6 +43,75 @@ class TestPortAvailability: sock.close() +class TestBackendAuthToken: + """Test internal Composer-to-Langflow backend authentication config.""" + + def test_backend_auth_member_config_adds_internal_header(self, mcp_service): + project_id = "backend-auth-test" + + config = mcp_service._build_backend_auth_member_config( + project_id, + "http://localhost:7860/api/v1/mcp/project/test/streamable", + legacy_sse_url="http://localhost:7860/api/v1/mcp/project/test/sse", + ) + + token = mcp_service.get_or_create_backend_auth_token(project_id) + assert mcp_service.validate_backend_auth_token(project_id, token) + assert [entry["type"] for entry in config] == ["http", "sse"] + assert config[0]["headers"] == {COMPOSER_BACKEND_AUTH_HEADER: token} + assert config[1]["headers"] == {COMPOSER_BACKEND_AUTH_HEADER: token} + + @pytest.mark.asyncio + async def test_oauth_startup_uses_config_path_for_backend_auth(self, mcp_service): + project_id = "backend-auth-process-test" + auth_config = { + "auth_type": "oauth", + "oauth_host": "localhost", + "oauth_port": "9000", + "oauth_server_url": "http://localhost:9000", + "oauth_client_id": "cid", + "oauth_client_secret": "csecret", # pragma: allowlist secret + "oauth_auth_url": "http://auth", + "oauth_token_url": "http://token", + } + + mock_settings = MagicMock() + mock_settings.settings.mcp_composer_version = "==0.1.0.8.10" + mock_process = MagicMock(pid=2222, poll=MagicMock(return_value=None)) + + with ( + patch("lfx.services.mcp_composer.service.get_settings_service", return_value=mock_settings), + patch.object( + mcp_service, + "_write_backend_auth_member_config", + return_value=Path("/tmp/mcp-composer-test-config.json"), + ) as mock_write_config, + patch("subprocess.Popen", return_value=mock_process) as mock_popen, + patch.object(mcp_service, "_is_port_available", return_value=False), + ): + await mcp_service._start_project_composer_process( + project_id=project_id, + host="localhost", + port=9000, + streamable_http_url="http://test/mcp", + auth_config=auth_config, + max_startup_checks=1, + startup_delay=0.01, + legacy_sse_url="http://test/sse", + ) + + cmd = mock_popen.call_args[0][0] + assert "--config_path" in cmd + assert "--endpoint" not in cmd + assert "--sse-url" not in cmd + assert COMPOSER_BACKEND_AUTH_HEADER not in cmd + mock_write_config.assert_called_once_with( + project_id, + "http://test/mcp", + legacy_sse_url="http://test/sse", + ) + + class TestKillProcessOnPort: """Test process killing functionality.""" From a46ae46f9f10c8c87a2bdad10e7260fc3c820b90 Mon Sep 17 00:00:00 2001 From: Eric Hare Date: Mon, 27 Apr 2026 19:02:40 -0700 Subject: [PATCH 3/3] feat: Add an Astra DB Data API Component (#12766) * feat: Add an Astra DB Data API Component * Run pre-commit hooks to format * Update astradb_data_api.py * Update component index * Merge branch 'release-1.10.0' into feat/astra-db-data-api-component * Updates * Update component_index.json * Update component_index.json --- docs/docs/Components/bundles-datastax.mdx | 99 +++ .../test_astradb_data_api_component.py | 714 +++++++++++++++ .../components/sideBarFolderButtons/index.tsx | 2 +- src/lfx/src/lfx/_assets/component_index.json | 810 +++++++++++++++++- .../src/lfx/components/datastax/__init__.py | 3 + .../components/datastax/astradb_data_api.py | 614 +++++++++++++ 6 files changed, 2239 insertions(+), 3 deletions(-) create mode 100644 src/backend/tests/unit/base/datastax/test_astradb_data_api_component.py create mode 100644 src/lfx/src/lfx/components/datastax/astradb_data_api.py diff --git a/docs/docs/Components/bundles-datastax.mdx b/docs/docs/Components/bundles-datastax.mdx index f365bf8b0e..8385a906cf 100644 --- a/docs/docs/Components/bundles-datastax.mdx +++ b/docs/docs/Components/bundles-datastax.mdx @@ -194,6 +194,105 @@ The output is a list of [`JSON`](/data-types#json) objects containing the query | Static Filters | Dict | Input parameter. Attribute-value pairs used to filter query results. | | Limit | String | Input parameter. The number of records to return. | +## Astra DB Data API + +The **Astra DB Data API** component runs document-level Data API operations — `find`, `find_one`, `insert_one`, `insert_many`, `update_one`, `update_many`, `delete_one`, `delete_many`, `count_documents`, and `estimated_document_count` — against a collection using the [`astrapy`](https://github.com/datastax/astrapy) Python SDK directly. + +Unlike the [**Astra DB** component](#astra-db), which wraps `AstraDBVectorStore` from `langchain-astradb` and is oriented around vector ingest/search, the **Astra DB Data API** component is a thin, direct pass-through to `astrapy`. No `langchain-astradb` code path is used at runtime, which makes it a good fit for: + +* Exact-match and compound metadata queries on collections that may or may not be vector-enabled. +* CRUD-style writes and updates (`$set`, `$inc`, `$push`, `$unset`, and other Data API operators). +* Document counts (bounded with `count_documents`, statistical with `estimated_document_count`). +* Wiring an agent tool that can *both* read and write to Astra DB. + +The component also provides an **astrapy-only** "Create new collection" dialog (equivalent to the one on the **Astra DB** component) so you can spin up a collection — with or without a vectorize integration — without installing `langchain-astradb`. + +:::note +The operation tab dynamically shows and hides only the inputs relevant to the selected operation, so the configuration surface stays minimal. For example, selecting **Insert One** hides all filter/projection/sort inputs and shows a single **Document** input. +::: + +### Astra DB Data API parameters + + + +| Name | Display Name | Info | +|------|--------------|------| +| token | Astra DB Application Token | Input parameter. An Astra application token with permission to access your database. | +| environment | Environment | Input parameter. The Astra DB environment, typically `prod`. | +| database_name | Database | Input parameter. The database the component will connect to. Supports in-app database creation. | +| api_endpoint | Astra DB API Endpoint | Input parameter. Auto-populated from the selected database. | +| keyspace | Keyspace | Input parameter. Defaults to `default_keyspace`. | +| collection_name | Collection | Input parameter. The target collection. Supports in-app collection creation via `astrapy` directly (no `langchain-astradb` dependency). | +| operation | Operation | Input parameter. One of `Find`, `Find One`, `Insert One`, `Insert Many`, `Update One`, `Update Many`, `Delete One`, `Delete Many`, `Count Documents`, `Estimated Count`. | +| filter_query | Filter | Input parameter. MongoDB-style filter document, e.g. `{"status": "active", "age": {"$gte": 18}}`. Used by find, update, delete, and count operations. | +| projection | Projection | Input parameter. Fields to include (`1`) or exclude (`0`), e.g. `{"name": 1, "email": 1}`. | +| sort | Sort | Input parameter. Sort specification. Use `$vector` or `$vectorize` for vector search, e.g. `{"$vectorize": "search query"}`. | +| limit | Limit | Input parameter. Maximum number of documents returned by `find`. Default: `100`. | +| skip | Skip | Input parameter. Number of documents to skip. | +| include_similarity | Include Similarity | Input parameter. Include the `$similarity` score on returned documents (vector searches only). | +| document | Document | Input parameter. Single JSON document for `Insert One`. | +| documents | Documents | Input parameter. List of JSON documents for `Insert Many`. | +| ordered | Ordered Insert | Input parameter. If `true`, `Insert Many` stops at the first error. Default: `false`. | +| update | Update | Input parameter. Update-operator document for `Update One` / `Update Many`, e.g. `{"$set": {"status": "archived"}}`. | +| upsert | Upsert | Input parameter. If `true`, inserts a new document when no match is found. Default: `false`. | +| upper_bound | Count Upper Bound | Input parameter. Maximum count `count_documents` will scan to. Default: `1000`. | + +### Astra DB Data API outputs + +| Output | Type | Description | +|--------|------|-------------| +| **Data** | `list[Data]` | List of JSON documents returned by the selected operation. Writes return a single-element list summarising the result (e.g. `inserted_ids`, `modified_count`, `deleted_count`). | +| **Table** | `Table` | The same result expressed as a tabular [`Table`](/data-types#table) for easy preview and downstream tabular processing. | +| **Raw Result** | `Data` | The raw `astrapy` result envelope as a single JSON object — useful for debugging or accessing fields like `matched_count` vs `modified_count`. | + +Because key inputs (`operation`, `filter_query`, `limit`, `document`, `documents`, `update`) expose **Tool Mode**, the component can also be connected to an **Agent** component as a tool. When used as a tool, the agent controls the selected operation and the payload while you retain control of the database, keyspace, and collection via the component's configuration. + +### Astra DB Data API examples + +
+Example: Insert documents + +1. In a flow, add the **Astra DB Data API** component. +2. Provide a token and select a database, keyspace, and collection. +3. Set **Operation** to **Insert Many**. +4. In **Documents**, provide a JSON list of documents, for example: + + ```json + [ + {"name": "Ada Lovelace", "role": "engineer"}, + {"name": "Grace Hopper", "role": "engineer"} + ] + ``` + +5. Run the component. The **Data** output contains a single summary object with `inserted_ids` and `inserted_count`. + +
+ +
+Example: Filter-based search + +1. Set **Operation** to **Find**. +2. In **Filter**, provide a MongoDB-style filter, for example: + + ```json + {"role": "engineer", "active": true} + ``` + +3. Optionally set **Projection** (e.g. `{"name": 1, "role": 1}`), **Sort**, and **Limit**. +4. Run the component. The **Data** output contains matching documents as a list of JSON objects; the **Table** output renders the same result as a [`Table`](/data-types#table). + +
+ +
+Example: Use as an agent tool + +1. In a flow with an **Agent** component, add the **Astra DB Data API** component. +2. Select the database, keyspace, and collection. +3. Toggle **Tool Mode** on the **Astra DB Data API** component, and connect its tool output to the agent's **Tools** input. +4. Leave **Operation**, **Filter**, **Limit**, **Document**, **Documents**, and **Update** unset — the agent will populate them per tool call based on the user's request. + +
+ ## Graph RAG The **Graph RAG** component uses an instance of [`GraphRetriever`](https://datastax.github.io/graph-rag/reference/langchain_graph_retriever/) for Graph RAG traversal enabling graph-based document retrieval in an Astra DB vector store. diff --git a/src/backend/tests/unit/base/datastax/test_astradb_data_api_component.py b/src/backend/tests/unit/base/datastax/test_astradb_data_api_component.py new file mode 100644 index 0000000000..9f9da5b43c --- /dev/null +++ b/src/backend/tests/unit/base/datastax/test_astradb_data_api_component.py @@ -0,0 +1,714 @@ +"""Unit tests for :class:`AstraDBDataAPIComponent`. + +These tests exercise every operation dispatch path and the astrapy-backed +collection-creation override without reaching out to Astra DB; ``astrapy`` +entry points are mocked. +""" + +from __future__ import annotations + +from collections import defaultdict +from unittest.mock import Mock, patch + +import pytest +from lfx.components.datastax.astradb_data_api import ( + ALL_OPERATIONS, + DEFAULT_COUNT_UPPER_BOUND, + OP_COUNT, + OP_DELETE_MANY, + OP_DELETE_ONE, + OP_ESTIMATED_COUNT, + OP_FIND, + OP_FIND_ONE, + OP_INSERT_MANY, + OP_INSERT_ONE, + OP_UPDATE_MANY, + OP_UPDATE_ONE, + OPERATION_FIELDS, + OPERATION_ICONS, + AstraDBDataAPIComponent, + _coerce_documents, + _stringify, +) +from lfx.schema.data import Data + +# ---------------------------------------------------------------------- +# Fixtures +# ---------------------------------------------------------------------- + + +@pytest.fixture +def component() -> AstraDBDataAPIComponent: + """Minimal component with required selector attrs pre-populated.""" + comp = AstraDBDataAPIComponent() + comp.token = "test_token" # noqa: S105 + comp.environment = "prod" + comp.database_name = "test_db" + comp.api_endpoint = "https://abcd1234-1234-1234-1234-1234567890ab.apps.astra.datastax.com" + comp.keyspace = "default_keyspace" + comp.collection_name = "test_coll" + comp.operation = OP_FIND + # Every operation-specific field starts unset. + comp.filter_query = {} + comp.projection = {} + comp.sort = {} + comp.limit = 0 + comp.skip = 0 + comp.include_similarity = False + comp.document = {} + comp.documents = [] + comp.ordered = False + comp.update = {} + comp.upsert = False + comp.upper_bound = DEFAULT_COUNT_UPPER_BOUND + comp.log = Mock() + return comp + + +@pytest.fixture +def mock_collection() -> Mock: + """A mocked astrapy.Collection with the operations we need.""" + return Mock() + + +@pytest.fixture +def component_with_collection(component, mock_collection): + """Component where ``_get_collection`` is patched to return ``mock_collection``.""" + component._get_collection = Mock(return_value=mock_collection) # type: ignore[method-assign] + return component + + +# ---------------------------------------------------------------------- +# Static / configuration tests +# ---------------------------------------------------------------------- + + +class TestStaticConfiguration: + """Tests that don't require a component instance.""" + + def test_all_operations_have_icons(self): + """Every dropdown option has an associated icon.""" + assert set(OPERATION_ICONS.keys()) == set(ALL_OPERATIONS) + + def test_all_operations_have_fields(self): + """Every dropdown option has a field-set definition (possibly empty).""" + assert set(OPERATION_FIELDS.keys()) == set(ALL_OPERATIONS) + + def test_operation_field_names_exist_on_inputs(self): + """Every field referenced in OPERATION_FIELDS is an actual declared input.""" + input_names = {inp.name for inp in AstraDBDataAPIComponent.inputs} + for op, fields in OPERATION_FIELDS.items(): + for field in fields: + assert field in input_names, f"Operation '{op}' references unknown input field '{field}'" + + def test_component_metadata(self): + """Component metadata is sensible.""" + assert AstraDBDataAPIComponent.display_name == "Astra DB Data API" + assert AstraDBDataAPIComponent.name == "AstraDBDataAPI" + assert AstraDBDataAPIComponent.icon == "AstraDB" + assert AstraDBDataAPIComponent.beta is True + + def test_outputs_defined(self): + """The three documented output sockets are present.""" + names = {o.name for o in AstraDBDataAPIComponent.outputs} + assert names == {"data", "dataframe", "raw"} + + def test_tool_mode_inputs(self): + """Key agent-controllable inputs expose ``tool_mode=True``.""" + tool_mode_names = {inp.name for inp in AstraDBDataAPIComponent.inputs if getattr(inp, "tool_mode", False)} + expected = {"operation", "filter_query", "limit", "document", "documents", "update"} + assert expected.issubset(tool_mode_names), ( + f"Missing tool_mode on agent-controllable inputs: {expected - tool_mode_names}" + ) + # Secrets and selector fields should never be marked as tool-controllable. + forbidden = {"token", "api_endpoint", "database_name", "collection_name", "keyspace"} + assert forbidden.isdisjoint(tool_mode_names), ( + f"Sensitive/selector fields must not be tool_mode: {forbidden & tool_mode_names}" + ) + + +# ---------------------------------------------------------------------- +# Helpers +# ---------------------------------------------------------------------- + + +class TestStringify: + def test_primitives_pass_through(self): + assert _stringify("abc") == "abc" + assert _stringify(1) == 1 + assert _stringify(1.5) == 1.5 + assert _stringify(True) is True # noqa: FBT003 + + def test_none_returns_none(self): + assert _stringify(None) is None + + def test_objects_stringified(self): + class Id: + def __str__(self): + return "object-id" + + assert _stringify(Id()) == "object-id" + + +class TestCoerceDocuments: + def test_empty_value(self): + assert _coerce_documents(None) == [] + assert _coerce_documents("") == [] + + def test_single_dict_wrapped(self): + result = _coerce_documents({"a": 1}) + assert result == [{"a": 1}] + + def test_list_passes_through(self): + docs = [{"a": 1}, {"b": 2}] + assert _coerce_documents(docs) == docs + + def test_rejects_non_dict_entries(self): + with pytest.raises(ValueError, match="non-dict entries"): + _coerce_documents([{"a": 1}, "not-a-dict"]) + + def test_rejects_wrong_type(self): + with pytest.raises(ValueError, match="Documents must be a list"): + _coerce_documents(42) + + +# ---------------------------------------------------------------------- +# UI: operation visibility toggling +# ---------------------------------------------------------------------- + + +class TestOperationVisibility: + def _make_build_config(self) -> dict: + """Build a config dict containing every toggle-able field.""" + fields = set().union(*OPERATION_FIELDS.values()) + return {field: {"show": True} for field in fields} + + def test_find_shows_query_fields(self): + config = self._make_build_config() + AstraDBDataAPIComponent._apply_operation_visibility(config, OP_FIND) + assert config["filter_query"]["show"] is True + assert config["limit"]["show"] is True + assert config["document"]["show"] is False + assert config["update"]["show"] is False + + def test_insert_one_shows_document(self): + config = self._make_build_config() + AstraDBDataAPIComponent._apply_operation_visibility(config, OP_INSERT_ONE) + assert config["document"]["show"] is True + assert config["documents"]["show"] is False + assert config["filter_query"]["show"] is False + + def test_insert_many_shows_documents_and_ordered(self): + config = self._make_build_config() + AstraDBDataAPIComponent._apply_operation_visibility(config, OP_INSERT_MANY) + assert config["documents"]["show"] is True + assert config["ordered"]["show"] is True + assert config["document"]["show"] is False + + def test_update_ops_show_filter_update_upsert(self): + for op in (OP_UPDATE_ONE, OP_UPDATE_MANY): + config = self._make_build_config() + AstraDBDataAPIComponent._apply_operation_visibility(config, op) + assert config["filter_query"]["show"] is True + assert config["update"]["show"] is True + assert config["upsert"]["show"] is True + assert config["document"]["show"] is False + + def test_count_shows_upper_bound(self): + config = self._make_build_config() + AstraDBDataAPIComponent._apply_operation_visibility(config, OP_COUNT) + assert config["filter_query"]["show"] is True + assert config["upper_bound"]["show"] is True + + def test_estimated_count_hides_everything(self): + config = self._make_build_config() + AstraDBDataAPIComponent._apply_operation_visibility(config, OP_ESTIMATED_COUNT) + for key, value in config.items(): + assert value["show"] is False, f"Field '{key}' should be hidden for Estimated Count" + + +# ---------------------------------------------------------------------- +# Operation dispatch +# ---------------------------------------------------------------------- + + +class TestFindOperation: + def test_find_returns_list_of_data(self, component_with_collection, mock_collection): + mock_collection.find.return_value = iter([{"_id": "1", "name": "Ada"}, {"_id": "2", "name": "Grace"}]) + component_with_collection.operation = OP_FIND + component_with_collection.filter_query = {"active": True} + component_with_collection.limit = 10 + + result = component_with_collection.run_operation() + + mock_collection.find.assert_called_once() + call_kwargs = mock_collection.find.call_args.kwargs + assert call_kwargs["filter"] == {"active": True} + assert call_kwargs["limit"] == 10 + + assert len(result) == 2 + assert all(isinstance(d, Data) for d in result) + assert result[0].data == {"_id": "1", "name": "Ada"} + + def test_find_omits_empty_options(self, component_with_collection, mock_collection): + mock_collection.find.return_value = iter([]) + component_with_collection.operation = OP_FIND + + component_with_collection.run_operation() + + kwargs = mock_collection.find.call_args.kwargs + assert "projection" not in kwargs + assert "sort" not in kwargs + assert "limit" not in kwargs + assert "skip" not in kwargs + assert "include_similarity" not in kwargs + + def test_find_respects_skip_projection_sort(self, component_with_collection, mock_collection): + mock_collection.find.return_value = iter([]) + component_with_collection.operation = OP_FIND + component_with_collection.projection = {"name": 1} + component_with_collection.sort = {"$vectorize": "hello"} + component_with_collection.skip = 5 + component_with_collection.limit = 20 + component_with_collection.include_similarity = True + + component_with_collection.run_operation() + + kwargs = mock_collection.find.call_args.kwargs + assert kwargs["projection"] == {"name": 1} + assert kwargs["sort"] == {"$vectorize": "hello"} + assert kwargs["skip"] == 5 + assert kwargs["limit"] == 20 + assert kwargs["include_similarity"] is True + + def test_find_raw_result(self, component_with_collection, mock_collection): + mock_collection.find.return_value = iter([{"a": 1}, {"a": 2}]) + component_with_collection.operation = OP_FIND + + raw = component_with_collection.raw_result() + + assert isinstance(raw, Data) + assert raw.data["count"] == 2 + assert raw.data["documents"] == [{"a": 1}, {"a": 2}] + + +class TestFindOneOperation: + def test_find_one_with_match(self, component_with_collection, mock_collection): + mock_collection.find_one.return_value = {"_id": "1", "name": "Ada"} + component_with_collection.operation = OP_FIND_ONE + component_with_collection.filter_query = {"_id": "1"} + + result = component_with_collection.run_operation() + + mock_collection.find_one.assert_called_once() + kwargs = mock_collection.find_one.call_args.kwargs + assert kwargs["filter"] == {"_id": "1"} + # ``limit`` / ``skip`` are not valid for find_one -- make sure we don't leak them. + assert "limit" not in kwargs + assert "skip" not in kwargs + assert len(result) == 1 + assert result[0].data == {"_id": "1", "name": "Ada"} + + def test_find_one_no_match_returns_empty(self, component_with_collection, mock_collection): + mock_collection.find_one.return_value = None + component_with_collection.operation = OP_FIND_ONE + component_with_collection.filter_query = {"_id": "missing"} + + assert component_with_collection.run_operation() == [] + + +class TestInsertOneOperation: + def test_insert_one_success(self, component_with_collection, mock_collection): + mock_result = Mock() + mock_result.inserted_id = "inserted-id-123" + mock_collection.insert_one.return_value = mock_result + + component_with_collection.operation = OP_INSERT_ONE + component_with_collection.document = {"name": "Ada"} + + result = component_with_collection.run_operation() + + mock_collection.insert_one.assert_called_once_with({"name": "Ada"}) + assert len(result) == 1 + assert result[0].data == {"inserted_id": "inserted-id-123"} + + def test_insert_one_rejects_non_dict(self, component_with_collection): + component_with_collection.operation = OP_INSERT_ONE + component_with_collection.document = "not-a-dict" + + with pytest.raises(ValueError, match="Astra DB Data API 'Insert One' failed"): + component_with_collection.run_operation() + + +class TestInsertManyOperation: + def test_insert_many_success(self, component_with_collection, mock_collection): + mock_result = Mock() + mock_result.inserted_ids = ["id-1", "id-2", "id-3"] + mock_collection.insert_many.return_value = mock_result + + component_with_collection.operation = OP_INSERT_MANY + component_with_collection.documents = [{"a": 1}, {"a": 2}, {"a": 3}] + component_with_collection.ordered = True + + result = component_with_collection.run_operation() + + mock_collection.insert_many.assert_called_once_with([{"a": 1}, {"a": 2}, {"a": 3}], ordered=True) + assert result[0].data == { + "inserted_ids": ["id-1", "id-2", "id-3"], + "inserted_count": 3, + } + + def test_insert_many_empty_rejected(self, component_with_collection): + component_with_collection.operation = OP_INSERT_MANY + component_with_collection.documents = [] + + with pytest.raises(ValueError, match="non-empty list"): + component_with_collection.run_operation() + + +class TestUpdateOperations: + def test_update_one_requires_filter(self, component_with_collection): + component_with_collection.operation = OP_UPDATE_ONE + component_with_collection.filter_query = {} + component_with_collection.update = {"$set": {"status": "x"}} + + with pytest.raises(ValueError, match="filter_query"): + component_with_collection.run_operation() + + def test_update_one_requires_update_doc(self, component_with_collection): + component_with_collection.operation = OP_UPDATE_ONE + component_with_collection.filter_query = {"_id": "1"} + component_with_collection.update = {} + + with pytest.raises(ValueError, match="update"): + component_with_collection.run_operation() + + def test_update_one_success(self, component_with_collection, mock_collection): + mock_result = Mock() + mock_result.matched_count = 1 + mock_result.modified_count = 1 + mock_result.upserted_id = None + mock_collection.update_one.return_value = mock_result + + component_with_collection.operation = OP_UPDATE_ONE + component_with_collection.filter_query = {"_id": "1"} + component_with_collection.update = {"$set": {"status": "active"}} + component_with_collection.upsert = True + + result = component_with_collection.run_operation() + + mock_collection.update_one.assert_called_once_with({"_id": "1"}, {"$set": {"status": "active"}}, upsert=True) + assert result[0].data["matched_count"] == 1 + assert result[0].data["modified_count"] == 1 + assert result[0].data["upserted_id"] is None + + def test_update_many_success(self, component_with_collection, mock_collection): + mock_result = Mock() + mock_result.matched_count = 5 + mock_result.modified_count = 5 + mock_result.upserted_id = None + mock_collection.update_many.return_value = mock_result + + component_with_collection.operation = OP_UPDATE_MANY + component_with_collection.filter_query = {"active": True} + component_with_collection.update = {"$inc": {"version": 1}} + + result = component_with_collection.run_operation() + + mock_collection.update_many.assert_called_once_with({"active": True}, {"$inc": {"version": 1}}, upsert=False) + assert result[0].data["modified_count"] == 5 + + +class TestDeleteOperations: + def test_delete_one(self, component_with_collection, mock_collection): + mock_result = Mock() + mock_result.deleted_count = 1 + mock_collection.delete_one.return_value = mock_result + + component_with_collection.operation = OP_DELETE_ONE + component_with_collection.filter_query = {"_id": "1"} + + result = component_with_collection.run_operation() + + mock_collection.delete_one.assert_called_once_with({"_id": "1"}) + assert result[0].data == {"deleted_count": 1} + + def test_delete_many(self, component_with_collection, mock_collection): + mock_result = Mock() + mock_result.deleted_count = 42 + mock_collection.delete_many.return_value = mock_result + + component_with_collection.operation = OP_DELETE_MANY + component_with_collection.filter_query = {"archived": True} + + result = component_with_collection.run_operation() + + mock_collection.delete_many.assert_called_once_with({"archived": True}) + assert result[0].data == {"deleted_count": 42} + + def test_delete_requires_filter(self, component_with_collection): + component_with_collection.operation = OP_DELETE_ONE + component_with_collection.filter_query = {} + + with pytest.raises(ValueError, match="filter_query"): + component_with_collection.run_operation() + + +class TestCountOperations: + def test_count_documents(self, component_with_collection, mock_collection): + mock_collection.count_documents.return_value = 17 + + component_with_collection.operation = OP_COUNT + component_with_collection.filter_query = {"active": True} + component_with_collection.upper_bound = 500 + + result = component_with_collection.run_operation() + + mock_collection.count_documents.assert_called_once_with({"active": True}, upper_bound=500) + assert result[0].data == {"count": 17, "upper_bound": 500} + + def test_estimated_count(self, component_with_collection, mock_collection): + mock_collection.estimated_document_count.return_value = 1_000_000 + + component_with_collection.operation = OP_ESTIMATED_COUNT + + result = component_with_collection.run_operation() + + mock_collection.estimated_document_count.assert_called_once_with() + assert result[0].data == {"estimated_count": 1_000_000} + + +# ---------------------------------------------------------------------- +# DataFrame output +# ---------------------------------------------------------------------- + + +class TestDataFrameOutput: + def test_as_dataframe_from_find(self, component_with_collection, mock_collection): + mock_collection.find.return_value = iter( + [ + {"name": "Ada", "age": 36}, + {"name": "Grace", "age": 85}, + ] + ) + component_with_collection.operation = OP_FIND + + df = component_with_collection.as_dataframe() + + assert len(df) == 2 + assert set(df.columns) >= {"name", "age"} + + def test_as_dataframe_empty(self, component_with_collection, mock_collection): + mock_collection.find.return_value = iter([]) + component_with_collection.operation = OP_FIND + + df = component_with_collection.as_dataframe() + + assert len(df) == 0 + + +# ---------------------------------------------------------------------- +# Error propagation +# ---------------------------------------------------------------------- + + +class TestErrorPropagation: + def test_unknown_operation_raises(self, component_with_collection): + component_with_collection.operation = "Destroy The Universe" + + with pytest.raises(ValueError, match="Unsupported operation"): + component_with_collection.run_operation() + + def test_astrapy_errors_wrapped(self, component_with_collection, mock_collection): + mock_collection.find.side_effect = RuntimeError("boom") + component_with_collection.operation = OP_FIND + + with pytest.raises(ValueError, match="Astra DB Data API 'Find' failed: boom"): + component_with_collection.run_operation() + + def test_missing_collection_name_raises(self, component): + component.collection_name = "" + # Bypass the patched _get_collection to exercise the real path. + with pytest.raises(ValueError, match="No collection selected"): + component._get_collection() + + +# ---------------------------------------------------------------------- +# update_build_config +# ---------------------------------------------------------------------- + + +class TestUpdateBuildConfig: + @pytest.mark.asyncio + async def test_operation_change_toggles_fields(self, component): + # Simulate the base behavior: reset_build_config if no token, so provide one. + build_config = { + "token": {"value": "tok"}, + "environment": {"value": "prod"}, + "database_name": {"value": "", "options": [], "options_metadata": [], "show": False}, + "api_endpoint": {"value": "", "options": []}, + "keyspace": {"value": "", "options": []}, + "collection_name": {"value": "", "options": [], "options_metadata": [], "show": False}, + "autodetect_collection": {"value": True}, + "operation": {"value": OP_FIND}, + # Fields managed by _apply_operation_visibility + "filter_query": {"show": True}, + "projection": {"show": True}, + "sort": {"show": True}, + "limit": {"show": True}, + "skip": {"show": True}, + "include_similarity": {"show": True}, + "document": {"show": True}, + "documents": {"show": True}, + "ordered": {"show": True}, + "update": {"show": True}, + "upsert": {"show": True}, + "upper_bound": {"show": True}, + } + + # Switching to Insert One should hide find-style fields. + with ( + patch.object( + AstraDBDataAPIComponent, + "get_database_list_static", + return_value={}, + ), + patch.object( + AstraDBDataAPIComponent, + "map_cloud_providers", + return_value={}, + ), + ): + # Swap to a non-managed field so we don't trigger the full base reset, + # then directly invoke the visibility helper via the operation branch. + result = await component.update_build_config(build_config, OP_INSERT_ONE, "operation") + + assert result["document"]["show"] is True + assert result["filter_query"]["show"] is False + assert result["documents"]["show"] is False + + +# ---------------------------------------------------------------------- +# astrapy-backed collection creation +# ---------------------------------------------------------------------- + + +class TestCreateCollectionAstrapyOnly: + """The override must not touch ``langchain-astradb``.""" + + @pytest.mark.asyncio + @patch("lfx.components.datastax.astradb_data_api.DataAPIClient") + async def test_bring_your_own_dimension(self, mock_client_cls): + mock_db = Mock() + mock_client_cls.return_value.get_database.return_value = mock_db + + await AstraDBDataAPIComponent.create_collection_api( + new_collection_name="c1", + token="t", # noqa: S106 + api_endpoint="https://x.apps.astra.datastax.com", + keyspace="ks", + dimension=1536, + ) + + mock_db.create_collection.assert_called_once() + args, kwargs = mock_db.create_collection.call_args + assert args[0] == "c1" + assert kwargs["keyspace"] == "ks" + definition = kwargs["definition"] + assert definition is not None + assert definition.vector is not None + assert definition.vector.dimension == 1536 + # Vectorize service must NOT be set when dimension-only. + assert definition.vector.service is None + + @pytest.mark.asyncio + @patch("lfx.components.datastax.astradb_data_api.DataAPIClient") + @patch.object(AstraDBDataAPIComponent, "get_vectorize_providers") + async def test_vectorize_provider(self, mock_providers, mock_client_cls): + mock_db = Mock() + mock_client_cls.return_value.get_database.return_value = mock_db + mock_providers.return_value = defaultdict( + list, + {"NVIDIA": ["nvidia", ["nv-embed-qa"]]}, + ) + + await AstraDBDataAPIComponent.create_collection_api( + new_collection_name="c2", + token="t", # noqa: S106 + api_endpoint="https://x.apps.astra.datastax.com", + keyspace="ks", + embedding_generation_provider="NVIDIA", + embedding_generation_model="nv-embed-qa", + ) + + kwargs = mock_db.create_collection.call_args.kwargs + definition = kwargs["definition"] + assert definition.vector.service is not None + assert definition.vector.service.provider == "nvidia" + assert definition.vector.service.model_name == "nv-embed-qa" + + @pytest.mark.asyncio + @patch("lfx.components.datastax.astradb_data_api.DataAPIClient") + async def test_plain_collection_no_vector(self, mock_client_cls): + """No dimension and no vectorize provider --> plain collection, no definition.""" + mock_db = Mock() + mock_client_cls.return_value.get_database.return_value = mock_db + + await AstraDBDataAPIComponent.create_collection_api( + new_collection_name="plain", + token="t", # noqa: S106 + api_endpoint="https://x.apps.astra.datastax.com", + keyspace="ks", + embedding_generation_provider="Bring your own", + ) + + kwargs = mock_db.create_collection.call_args.kwargs + assert kwargs["definition"] is None + + @pytest.mark.asyncio + async def test_missing_name_rejected(self): + with pytest.raises(ValueError, match="Collection name is required"): + await AstraDBDataAPIComponent.create_collection_api( + new_collection_name="", + token="t", # noqa: S106 + api_endpoint="https://x.apps.astra.datastax.com", + ) + + @pytest.mark.asyncio + @patch("lfx.components.datastax.astradb_data_api.DataAPIClient") + @patch.object(AstraDBDataAPIComponent, "get_vectorize_providers") + async def test_unknown_provider_rejected(self, mock_providers, mock_client_cls): + mock_client_cls.return_value.get_database.return_value = Mock() + mock_providers.return_value = defaultdict(list) + + with pytest.raises(ValueError, match="Unknown embedding provider"): + await AstraDBDataAPIComponent.create_collection_api( + new_collection_name="c3", + token="t", # noqa: S106 + api_endpoint="https://x.apps.astra.datastax.com", + embedding_generation_provider="NotAProvider", + embedding_generation_model="m", + ) + + @pytest.mark.asyncio + @patch("lfx.components.datastax.astradb_data_api.DataAPIClient") + async def test_no_langchain_astradb_import(self, mock_client_cls, monkeypatch): + """The override must not import ``langchain_astradb``. + + We simulate that module being unavailable; the call must still succeed. + """ + import sys + + monkeypatch.setitem(sys.modules, "langchain_astradb", None) + + mock_db = Mock() + mock_client_cls.return_value.get_database.return_value = mock_db + + await AstraDBDataAPIComponent.create_collection_api( + new_collection_name="no-lc", + token="t", # noqa: S106 + api_endpoint="https://x.apps.astra.datastax.com", + dimension=512, + ) + + mock_db.create_collection.assert_called_once() diff --git a/src/frontend/src/components/core/folderSidebarComponent/components/sideBarFolderButtons/index.tsx b/src/frontend/src/components/core/folderSidebarComponent/components/sideBarFolderButtons/index.tsx index e6fd188c95..0a808bff50 100644 --- a/src/frontend/src/components/core/folderSidebarComponent/components/sideBarFolderButtons/index.tsx +++ b/src/frontend/src/components/core/folderSidebarComponent/components/sideBarFolderButtons/index.tsx @@ -1,5 +1,6 @@ import { useIsFetching, useIsMutating } from "@tanstack/react-query"; import { useEffect, useRef, useState } from "react"; +import { useTranslation } from "react-i18next"; import { useLocation, useParams } from "react-router-dom"; import ForwardedIconComponent from "@/components/common/genericIconComponent"; import { @@ -13,7 +14,6 @@ import { SidebarMenuButton, SidebarMenuItem, } from "@/components/ui/sidebar"; -import { useTranslation } from "react-i18next"; import { useUpdateUser } from "@/controllers/API/queries/auth"; import { usePatchFolders, diff --git a/src/lfx/src/lfx/_assets/component_index.json b/src/lfx/src/lfx/_assets/component_index.json index 932436d2db..b334e07447 100644 --- a/src/lfx/src/lfx/_assets/component_index.json +++ b/src/lfx/src/lfx/_assets/component_index.json @@ -61047,6 +61047,812 @@ }, "tool_mode": false }, + "AstraDBDataAPI": { + "base_classes": [ + "JSON", + "Table" + ], + "beta": true, + "conditional_paths": [], + "custom_fields": {}, + "description": "Run Data API operations (find, filter, insert, update, delete, count) against an Astra DB collection using astrapy directly.", + "display_name": "Astra DB Data API", + "documentation": "https://docs.datastax.com/en/astra-db-serverless/api-reference/overview.html", + "edited": false, + "field_order": [ + "token", + "environment", + "database_name", + "api_endpoint", + "keyspace", + "collection_name", + "autodetect_collection", + "operation", + "filter_query", + "projection", + "sort", + "limit", + "skip", + "include_similarity", + "document", + "documents", + "ordered", + "update", + "upsert", + "upper_bound" + ], + "frozen": false, + "icon": "AstraDB", + "legacy": false, + "metadata": { + "code_hash": "79efc628e485", + "dependencies": { + "dependencies": [ + { + "name": "astrapy", + "version": "2.2.1" + }, + { + "name": "lfx", + "version": null + } + ], + "total_dependencies": 2 + }, + "module": "lfx.components.datastax.astradb_data_api.AstraDBDataAPIComponent" + }, + "minimized": false, + "output_types": [], + "outputs": [ + { + "allows_loop": false, + "cache": true, + "display_name": "Data", + "group_outputs": false, + "method": "run_operation", + "name": "data", + "selected": "JSON", + "tool_mode": true, + "types": [ + "JSON" + ], + "value": "__UNDEFINED__" + }, + { + "allows_loop": false, + "cache": true, + "display_name": "Table", + "group_outputs": false, + "method": "as_dataframe", + "name": "dataframe", + "selected": "Table", + "tool_mode": true, + "types": [ + "Table" + ], + "value": "__UNDEFINED__" + }, + { + "allows_loop": false, + "cache": true, + "display_name": "Raw Result", + "group_outputs": false, + "method": "raw_result", + "name": "raw", + "selected": "JSON", + "tool_mode": true, + "types": [ + "JSON" + ], + "value": "__UNDEFINED__" + } + ], + "pinned": false, + "template": { + "_type": "Component", + "api_endpoint": { + "_input_type": "DropdownInput", + "advanced": true, + "combobox": false, + "dialog_inputs": {}, + "display_name": "Astra DB API Endpoint", + "dynamic": false, + "external_options": {}, + "info": "The API Endpoint for the Astra DB instance. Supercedes database selection.", + "name": "api_endpoint", + "options": [], + "options_metadata": [], + "override_skip": false, + "placeholder": "", + "required": false, + "show": true, + "title_case": false, + "toggle": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "str", + "value": "" + }, + "autodetect_collection": { + "_input_type": "BoolInput", + "advanced": true, + "display_name": "Autodetect Collection", + "dynamic": false, + "info": "Boolean flag to determine whether to autodetect the collection.", + "list": false, + "list_add_label": "Add More", + "name": "autodetect_collection", + "override_skip": false, + "placeholder": "", + "required": false, + "show": true, + "title_case": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "bool", + "value": true + }, + "code": { + "advanced": true, + "dynamic": true, + "fileTypes": [], + "file_path": "", + "info": "", + "list": false, + "load_from_db": false, + "multiline": true, + "name": "code", + "password": false, + "placeholder": "", + "required": true, + "show": true, + "title_case": false, + "type": "code", + "value": "\"\"\"Astra DB Data API component.\n\nA thin, modern Langflow component that exposes the full document-based\nData API surface of DataStax Astra DB using **only** the ``astrapy`` SDK.\n\nThe component inherits :class:`lfx.base.datastax.astradb_base.AstraDBBaseComponent`\nto reuse the polished database / collection / keyspace selector UI (including\nthe in-app \"Create new database\" and \"Create new collection\" dialogs).\nUnlike :class:`AstraDBVectorStoreComponent`, no ``langchain-astradb`` code path\nis used at runtime --- every operation goes through ``astrapy`` directly.\n\nSupported operations (selectable via the operation tab):\n\n* Find --- ``collection.find``\n* Find One --- ``collection.find_one``\n* Insert One --- ``collection.insert_one``\n* Insert Many --- ``collection.insert_many``\n* Update One --- ``collection.update_one``\n* Update Many --- ``collection.update_many``\n* Delete One --- ``collection.delete_one``\n* Delete Many --- ``collection.delete_many``\n* Count Documents --- ``collection.count_documents``\n* Estimated Count --- ``collection.estimated_document_count``\n\"\"\"\n\nfrom __future__ import annotations\n\nfrom typing import Any\n\nfrom astrapy import Collection, DataAPIClient, Database\nfrom astrapy.info import (\n CollectionDefinition,\n CollectionVectorOptions,\n VectorServiceOptions,\n)\n\nfrom lfx.base.datastax.astradb_base import AstraDBBaseComponent\nfrom lfx.io import (\n BoolInput,\n DropdownInput,\n IntInput,\n NestedDictInput,\n Output,\n)\nfrom lfx.log.logger import logger\nfrom lfx.schema.data import Data\nfrom lfx.schema.dataframe import DataFrame\n\n# Operation option constants -- kept as module-level constants so the UI\n# tab values and the dispatcher stay in sync with a single source of truth.\nOP_FIND = \"Find\"\nOP_FIND_ONE = \"Find One\"\nOP_INSERT_ONE = \"Insert One\"\nOP_INSERT_MANY = \"Insert Many\"\nOP_UPDATE_ONE = \"Update One\"\nOP_UPDATE_MANY = \"Update Many\"\nOP_DELETE_ONE = \"Delete One\"\nOP_DELETE_MANY = \"Delete Many\"\nOP_COUNT = \"Count Documents\"\nOP_ESTIMATED_COUNT = \"Estimated Count\"\n\nALL_OPERATIONS: tuple[str, ...] = (\n OP_FIND,\n OP_FIND_ONE,\n OP_INSERT_ONE,\n OP_INSERT_MANY,\n OP_UPDATE_ONE,\n OP_UPDATE_MANY,\n OP_DELETE_ONE,\n OP_DELETE_MANY,\n OP_COUNT,\n OP_ESTIMATED_COUNT,\n)\n\n# Per-operation icons surfaced in the operation dropdown. Icon names mirror\n# those used elsewhere in Langflow for a consistent look.\nOPERATION_ICONS: dict[str, str] = {\n OP_FIND: \"Search\",\n OP_FIND_ONE: \"SearchCheck\",\n OP_INSERT_ONE: \"Plus\",\n OP_INSERT_MANY: \"CopyPlus\",\n OP_UPDATE_ONE: \"Pencil\",\n OP_UPDATE_MANY: \"PencilLine\",\n OP_DELETE_ONE: \"Trash\",\n OP_DELETE_MANY: \"Trash2\",\n OP_COUNT: \"Hash\",\n OP_ESTIMATED_COUNT: \"Sigma\",\n}\n\n# Groups of inputs that each operation needs. Centralising this makes the\n# dynamic show/hide logic trivial to audit and extend.\nOPERATION_FIELDS: dict[str, set[str]] = {\n OP_FIND: {\"filter_query\", \"projection\", \"sort\", \"limit\", \"skip\", \"include_similarity\"},\n OP_FIND_ONE: {\"filter_query\", \"projection\", \"sort\", \"include_similarity\"},\n OP_INSERT_ONE: {\"document\"},\n OP_INSERT_MANY: {\"documents\", \"ordered\"},\n OP_UPDATE_ONE: {\"filter_query\", \"update\", \"upsert\"},\n OP_UPDATE_MANY: {\"filter_query\", \"update\", \"upsert\"},\n OP_DELETE_ONE: {\"filter_query\"},\n OP_DELETE_MANY: {\"filter_query\"},\n OP_COUNT: {\"filter_query\", \"upper_bound\"},\n OP_ESTIMATED_COUNT: set(),\n}\n\n# Default count upper bound (``count_documents`` in astrapy requires one).\nDEFAULT_COUNT_UPPER_BOUND = 1000\n# Default find limit -- guard against unbounded scans in the UI path.\nDEFAULT_FIND_LIMIT = 100\n\n# Fields specific to the operation that we toggle ``show`` on.\n_OPERATION_TOGGLE_FIELDS: tuple[str, ...] = (\n \"filter_query\",\n \"projection\",\n \"sort\",\n \"limit\",\n \"skip\",\n \"include_similarity\",\n \"document\",\n \"documents\",\n \"update\",\n \"upsert\",\n \"ordered\",\n \"upper_bound\",\n)\n\n\nclass AstraDBDataAPIComponent(AstraDBBaseComponent):\n \"\"\"Direct ``astrapy`` Data API access for Astra DB collections.\n\n Inherits the standard Astra DB selector UI (database / keyspace /\n collection, plus the create-new dialogs) from\n :class:`AstraDBBaseComponent` and adds a single operation tab that\n drives a minimal, operation-specific set of inputs.\n \"\"\"\n\n display_name: str = \"Astra DB Data API\"\n description: str = (\n \"Run Data API operations (find, filter, insert, update, delete, \"\n \"count) against an Astra DB collection using astrapy directly.\"\n )\n documentation: str = \"https://docs.datastax.com/en/astra-db-serverless/api-reference/overview.html\"\n icon: str = \"AstraDB\"\n name: str = \"AstraDBDataAPI\"\n beta: bool = True\n\n inputs = [\n *AstraDBBaseComponent.inputs,\n DropdownInput(\n name=\"operation\",\n display_name=\"Operation\",\n info=\"Data API operation to run against the selected collection.\",\n options=list(ALL_OPERATIONS),\n options_metadata=[{\"icon\": OPERATION_ICONS[op]} for op in ALL_OPERATIONS],\n value=OP_FIND,\n real_time_refresh=True,\n combobox=False,\n tool_mode=True,\n ),\n # -- Query / projection / sort ---------------------------------\n NestedDictInput(\n name=\"filter_query\",\n display_name=\"Filter\",\n info=(\n \"MongoDB-style filter, e.g. \"\n '{\"status\": \"active\", \"age\": {\"$gte\": 18}}. '\n \"Leave empty to match all documents (not recommended for \"\n \"delete/update operations).\"\n ),\n advanced=False,\n tool_mode=True,\n ),\n NestedDictInput(\n name=\"projection\",\n display_name=\"Projection\",\n info=('Fields to include (``1``/``true``) or exclude (``0``/``false``), e.g. {\"name\": 1, \"email\": 1}.'),\n advanced=True,\n ),\n NestedDictInput(\n name=\"sort\",\n display_name=\"Sort\",\n info=(\n \"Sort specification -- ``1`` ascending, ``-1`` descending. \"\n \"Use ``$vector`` or ``$vectorize`` for vector search, e.g. \"\n '{\"$vectorize\": \"search query\"}.'\n ),\n advanced=True,\n ),\n IntInput(\n name=\"limit\",\n display_name=\"Limit\",\n info=\"Maximum number of documents to return.\",\n value=DEFAULT_FIND_LIMIT,\n advanced=False,\n tool_mode=True,\n ),\n IntInput(\n name=\"skip\",\n display_name=\"Skip\",\n info=\"Number of documents to skip.\",\n value=0,\n advanced=True,\n ),\n BoolInput(\n name=\"include_similarity\",\n display_name=\"Include Similarity\",\n info=\"Include the ``$similarity`` score on each returned document (vector searches only).\",\n value=False,\n advanced=True,\n ),\n # -- Writes -----------------------------------------------------\n NestedDictInput(\n name=\"document\",\n display_name=\"Document\",\n info='Single document to insert, e.g. {\"name\": \"Ada\", \"age\": 36}.',\n show=False,\n tool_mode=True,\n ),\n NestedDictInput(\n name=\"documents\",\n display_name=\"Documents\",\n info=\"List of documents to insert. Provide a JSON list, e.g. [{...}, {...}].\",\n show=False,\n tool_mode=True,\n ),\n BoolInput(\n name=\"ordered\",\n display_name=\"Ordered Insert\",\n info=\"If enabled, inserts stop at the first error; otherwise errors are collected.\",\n value=False,\n show=False,\n advanced=True,\n ),\n NestedDictInput(\n name=\"update\",\n display_name=\"Update\",\n info=(\n 'Update operator document, e.g. {\"$set\": {\"status\": \"archived\"}}. '\n \"Must use update operators (``$set``, ``$inc``, ``$push`` ...).\"\n ),\n show=False,\n tool_mode=True,\n ),\n BoolInput(\n name=\"upsert\",\n display_name=\"Upsert\",\n info=\"Insert a new document if no match is found.\",\n value=False,\n show=False,\n advanced=True,\n ),\n # -- Count ------------------------------------------------------\n IntInput(\n name=\"upper_bound\",\n display_name=\"Count Upper Bound\",\n info=\"Maximum count astrapy will scan to for ``Count Documents``.\",\n value=DEFAULT_COUNT_UPPER_BOUND,\n show=False,\n advanced=True,\n ),\n ]\n\n outputs = [\n Output(display_name=\"Data\", name=\"data\", method=\"run_operation\"),\n Output(display_name=\"Table\", name=\"dataframe\", method=\"as_dataframe\"),\n Output(display_name=\"Raw Result\", name=\"raw\", method=\"raw_result\"),\n ]\n\n # ----------------------------------------------------------------------\n # Override collection creation to use astrapy *directly*.\n # ----------------------------------------------------------------------\n\n @classmethod\n async def create_collection_api( # type: ignore[override]\n cls,\n new_collection_name: str,\n token: str,\n api_endpoint: str,\n environment: str | None = None,\n keyspace: str | None = None,\n dimension: int | None = None,\n embedding_generation_provider: str | None = None,\n embedding_generation_model: str | None = None,\n ) -> Collection:\n \"\"\"Create a new Astra DB collection using astrapy only.\n\n The base class implementation in :class:`AstraDBBaseComponent` uses\n ``_AstraDBCollectionEnvironment`` from ``langchain-astradb``. This\n override instead calls ``astrapy``'s :meth:`Database.create_collection`\n directly, building a :class:`CollectionDefinition` when\n server-side vectorize options are requested.\n\n Args:\n new_collection_name: Required collection name.\n token: Astra DB application token.\n api_endpoint: Full Astra DB API endpoint for the target database.\n environment: Astra environment (``prod`` / ``test`` / ``dev``).\n keyspace: Keyspace to create the collection in.\n dimension: Embedding dimension if using \"Bring your own\" embeddings.\n embedding_generation_provider: Display name of a configured\n vectorize provider, or ``\"Bring your own\"``.\n embedding_generation_model: Model name for the selected provider.\n\n Returns:\n The newly created :class:`astrapy.Collection`.\n\n Raises:\n ValueError: If ``new_collection_name`` is empty or if an\n unsupported provider configuration is supplied.\n \"\"\"\n if not new_collection_name:\n msg = \"Collection name is required to create a new collection.\"\n raise ValueError(msg)\n\n env = cls.get_environment(environment)\n client = DataAPIClient(environment=env)\n database = client.get_database(api_endpoint, token=token, keyspace=keyspace)\n\n # Build the vector options -- either bring-your-own (dimension only)\n # or a server-side vectorize provider.\n vector_options: CollectionVectorOptions | None = None\n if dimension:\n vector_options = CollectionVectorOptions(dimension=dimension)\n elif embedding_generation_provider and embedding_generation_provider != \"Bring your own\":\n providers = cls.get_vectorize_providers(token=token, environment=env, api_endpoint=api_endpoint)\n provider_key = providers.get(embedding_generation_provider, [None, []])[0]\n if provider_key is None:\n msg = (\n f\"Unknown embedding provider '{embedding_generation_provider}'. \"\n \"Pass a dimension for a bring-your-own collection, or choose a \"\n \"configured vectorize provider.\"\n )\n raise ValueError(msg)\n vector_options = CollectionVectorOptions(\n service=VectorServiceOptions(\n provider=provider_key,\n model_name=embedding_generation_model,\n ),\n )\n\n definition = CollectionDefinition(vector=vector_options) if vector_options else None\n return database.create_collection(\n new_collection_name,\n definition=definition,\n keyspace=keyspace,\n )\n\n # ----------------------------------------------------------------------\n # Dynamic UI -- show/hide operation-specific fields.\n # ----------------------------------------------------------------------\n\n async def update_build_config(\n self,\n build_config: dict,\n field_value: str | dict,\n field_name: str | None = None,\n ) -> dict:\n \"\"\"Keep the base Astra DB selector behavior and toggle operation fields.\"\"\"\n build_config = await super().update_build_config(build_config, field_value, field_name)\n\n # On operation change (or first render), toggle visibility of\n # operation-specific inputs.\n if field_name == \"operation\" or field_name is None:\n op_value = field_value if field_name == \"operation\" else build_config.get(\"operation\", {}).get(\"value\")\n self._apply_operation_visibility(build_config, op_value or OP_FIND)\n\n return build_config\n\n @classmethod\n def _apply_operation_visibility(cls, build_config: dict, operation: str) -> None:\n \"\"\"Show/hide operation-specific fields based on the selected operation.\"\"\"\n active = OPERATION_FIELDS.get(operation, set())\n for field in _OPERATION_TOGGLE_FIELDS:\n if field in build_config:\n build_config[field][\"show\"] = field in active\n\n # ----------------------------------------------------------------------\n # Collection accessor\n # ----------------------------------------------------------------------\n\n def _get_collection(self) -> Collection:\n \"\"\"Return the resolved astrapy :class:`Collection` for this component.\"\"\"\n if not self.collection_name:\n msg = \"No collection selected. Choose or create a collection first.\"\n raise ValueError(msg)\n\n database: Database = self.get_database_object(api_endpoint=self.get_api_endpoint())\n return database.get_collection(self.collection_name, keyspace=self.get_keyspace())\n\n # ----------------------------------------------------------------------\n # Operation dispatch\n # ----------------------------------------------------------------------\n\n def run_operation(self) -> list[Data]:\n \"\"\"Execute the selected operation and return results as ``list[Data]``.\n\n For write / count operations where a ``list[Data]`` isn't semantically\n meaningful, a single-element list wrapping a summary :class:`Data`\n object is returned so the output socket stays consistent.\n \"\"\"\n result = self._dispatch()\n self.status = result\n return result\n\n def raw_result(self) -> Data:\n \"\"\"Return the raw operation result as a single :class:`Data` object.\n\n Useful for inspecting the full ``astrapy`` response envelope\n (``inserted_ids``, ``matched_count``, ``modified_count`` ...).\n \"\"\"\n raw = self._dispatch(raw=True)\n data = Data(data=raw if isinstance(raw, dict) else {\"result\": raw})\n self.status = data\n return data\n\n def as_dataframe(self) -> DataFrame:\n \"\"\"Return results as a :class:`DataFrame` for easy tabular display.\"\"\"\n rows = self.run_operation()\n return DataFrame([row.data for row in rows if isinstance(row, Data)])\n\n # ----------------------------------------------------------------------\n # Internal dispatcher\n # ----------------------------------------------------------------------\n\n def _dispatch(self, *, raw: bool = False) -> Any:\n operation = getattr(self, \"operation\", OP_FIND) or OP_FIND\n collection = self._get_collection()\n\n handlers = {\n OP_FIND: self._op_find,\n OP_FIND_ONE: self._op_find_one,\n OP_INSERT_ONE: self._op_insert_one,\n OP_INSERT_MANY: self._op_insert_many,\n OP_UPDATE_ONE: self._op_update_one,\n OP_UPDATE_MANY: self._op_update_many,\n OP_DELETE_ONE: self._op_delete_one,\n OP_DELETE_MANY: self._op_delete_many,\n OP_COUNT: self._op_count,\n OP_ESTIMATED_COUNT: self._op_estimated_count,\n }\n\n handler = handlers.get(operation)\n if handler is None:\n msg = f\"Unsupported operation '{operation}'.\"\n raise ValueError(msg)\n\n try:\n return handler(collection, raw=raw)\n except Exception as exc:\n logger.exception(\"Astra DB Data API operation failed: %s\", operation)\n msg = f\"Astra DB Data API '{operation}' failed: {exc}\"\n raise ValueError(msg) from exc\n\n # -- Handlers -----------------------------------------------------\n\n def _op_find(self, collection: Collection, *, raw: bool) -> Any:\n find_kwargs = self._build_find_kwargs()\n cursor = collection.find(**find_kwargs)\n docs = list(cursor)\n if raw:\n return {\"count\": len(docs), \"documents\": docs}\n return [Data(data=doc) for doc in docs]\n\n def _op_find_one(self, collection: Collection, *, raw: bool) -> Any:\n kwargs = self._build_find_kwargs(one=True)\n doc = collection.find_one(**kwargs)\n if raw:\n return {\"document\": doc}\n return [Data(data=doc)] if doc else []\n\n def _op_insert_one(self, collection: Collection, *, raw: bool) -> Any:\n document = self._require_mapping(\"document\", self.document, allow_empty=False)\n result = collection.insert_one(document)\n payload = {\"inserted_id\": _stringify(result.inserted_id)}\n if raw:\n return payload\n return [Data(data=payload)]\n\n def _op_insert_many(self, collection: Collection, *, raw: bool) -> Any:\n documents = _coerce_documents(self.documents)\n if not documents:\n msg = \"Insert Many requires a non-empty list of documents.\"\n raise ValueError(msg)\n result = collection.insert_many(documents, ordered=bool(self.ordered))\n payload = {\n \"inserted_ids\": [_stringify(_id) for _id in result.inserted_ids],\n \"inserted_count\": len(result.inserted_ids),\n }\n if raw:\n return payload\n return [Data(data=payload)]\n\n def _op_update_one(self, collection: Collection, *, raw: bool) -> Any:\n filter_q = self._require_mapping(\"filter_query\", self.filter_query, allow_empty=False)\n update_doc = self._require_mapping(\"update\", self.update, allow_empty=False)\n result = collection.update_one(filter_q, update_doc, upsert=bool(self.upsert))\n return self._format_update_result(result, raw=raw)\n\n def _op_update_many(self, collection: Collection, *, raw: bool) -> Any:\n filter_q = self._require_mapping(\"filter_query\", self.filter_query, allow_empty=False)\n update_doc = self._require_mapping(\"update\", self.update, allow_empty=False)\n result = collection.update_many(filter_q, update_doc, upsert=bool(self.upsert))\n return self._format_update_result(result, raw=raw)\n\n def _op_delete_one(self, collection: Collection, *, raw: bool) -> Any:\n filter_q = self._require_mapping(\"filter_query\", self.filter_query, allow_empty=False)\n result = collection.delete_one(filter_q)\n payload = {\"deleted_count\": result.deleted_count}\n if raw:\n return payload\n return [Data(data=payload)]\n\n def _op_delete_many(self, collection: Collection, *, raw: bool) -> Any:\n filter_q = self._require_mapping(\"filter_query\", self.filter_query, allow_empty=False)\n result = collection.delete_many(filter_q)\n payload = {\"deleted_count\": result.deleted_count}\n if raw:\n return payload\n return [Data(data=payload)]\n\n def _op_count(self, collection: Collection, *, raw: bool) -> Any:\n filter_q = self.filter_query or {}\n upper_bound = int(self.upper_bound or DEFAULT_COUNT_UPPER_BOUND)\n count = collection.count_documents(filter_q, upper_bound=upper_bound)\n payload = {\"count\": count, \"upper_bound\": upper_bound}\n if raw:\n return payload\n return [Data(data=payload)]\n\n def _op_estimated_count(self, collection: Collection, *, raw: bool) -> Any:\n count = collection.estimated_document_count()\n payload = {\"estimated_count\": count}\n if raw:\n return payload\n return [Data(data=payload)]\n\n # -- Helpers ------------------------------------------------------\n\n def _build_find_kwargs(self, *, one: bool = False) -> dict[str, Any]:\n \"\"\"Build kwargs for ``find`` / ``find_one``, omitting empty values.\"\"\"\n kwargs: dict[str, Any] = {\"filter\": self.filter_query or {}}\n\n if self.projection:\n kwargs[\"projection\"] = self.projection\n if self.sort:\n kwargs[\"sort\"] = self.sort\n if self.include_similarity:\n kwargs[\"include_similarity\"] = True\n\n if not one:\n limit = int(self.limit) if self.limit else 0\n if limit > 0:\n kwargs[\"limit\"] = limit\n skip = int(self.skip) if self.skip else 0\n if skip > 0:\n kwargs[\"skip\"] = skip\n return kwargs\n\n @staticmethod\n def _format_update_result(result: Any, *, raw: bool) -> Any:\n payload = {\n \"matched_count\": getattr(result, \"matched_count\", None),\n \"modified_count\": getattr(result, \"modified_count\", None),\n \"upserted_id\": _stringify(getattr(result, \"upserted_id\", None)),\n }\n if raw:\n return payload\n return [Data(data=payload)]\n\n @staticmethod\n def _require_mapping(field_name: str, value: Any, *, allow_empty: bool = True) -> dict:\n \"\"\"Validate that a field contains a mapping; raise a clear error otherwise.\"\"\"\n if value is None or value == \"\":\n if allow_empty:\n return {}\n msg = f\"Field '{field_name}' is required for this operation.\"\n raise ValueError(msg)\n if not isinstance(value, dict):\n msg = f\"Field '{field_name}' must be a JSON object, got {type(value).__name__}.\"\n raise ValueError(msg) # noqa: TRY004\n if not value and not allow_empty:\n msg = f\"Field '{field_name}' must not be empty for this operation.\"\n raise ValueError(msg)\n return value\n\n\n# ----------------------------------------------------------------------\n# Module-level helpers (kept outside the class for easy unit-testing)\n# ----------------------------------------------------------------------\n\n\ndef _stringify(value: Any) -> Any:\n \"\"\"Best-effort stringification of Data API id types for JSON-serialisable output.\"\"\"\n if value is None:\n return None\n if isinstance(value, str | int | float | bool):\n return value\n return str(value)\n\n\ndef _coerce_documents(value: Any) -> list[dict[str, Any]]:\n \"\"\"Accept either a ``list[dict]`` or a single ``dict`` and normalise to a list.\"\"\"\n if value is None or value == \"\":\n return []\n if isinstance(value, dict):\n # Singleton -- wrap into a list for convenience.\n return [value]\n if isinstance(value, list):\n bad = [i for i, item in enumerate(value) if not isinstance(item, dict)]\n if bad:\n msg = f\"Documents list contains non-dict entries at index {bad[0]}.\"\n raise ValueError(msg)\n return value\n msg = f\"Documents must be a list of JSON objects, got {type(value).__name__}.\"\n raise ValueError(msg)\n" + }, + "collection_name": { + "_input_type": "DropdownInput", + "advanced": false, + "combobox": true, + "dialog_inputs": { + "fields": { + "data": { + "node": { + "description": "Please allow several seconds for creation to complete.", + "display_name": "Create new collection", + "field_order": [ + "01_new_collection_name", + "02_embedding_generation_provider", + "03_embedding_generation_model", + "04_dimension" + ], + "name": "create_collection", + "template": { + "01_new_collection_name": { + "_input_type": "StrInput", + "advanced": false, + "display_name": "Name", + "dynamic": false, + "info": "Name of the new collection to create in Astra DB.", + "list": false, + "list_add_label": "Add More", + "load_from_db": false, + "name": "new_collection_name", + "override_skip": false, + "placeholder": "", + "required": true, + "show": true, + "title_case": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": false, + "type": "str", + "value": "" + }, + "02_embedding_generation_provider": { + "_input_type": "DropdownInput", + "advanced": false, + "combobox": false, + "dialog_inputs": {}, + "display_name": "Embedding generation method", + "dynamic": false, + "external_options": {}, + "helper_text": "To create collections with more embedding provider options, go to your database in Astra DB", + "info": "Provider to use for generating embeddings.", + "name": "embedding_generation_provider", + "options": [], + "options_metadata": [], + "override_skip": false, + "placeholder": "", + "real_time_refresh": true, + "required": true, + "show": true, + "title_case": false, + "toggle": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "str", + "value": "" + }, + "03_embedding_generation_model": { + "_input_type": "DropdownInput", + "advanced": false, + "combobox": false, + "dialog_inputs": {}, + "display_name": "Embedding model", + "dynamic": false, + "external_options": {}, + "info": "Model to use for generating embeddings.", + "name": "embedding_generation_model", + "options": [], + "options_metadata": [], + "override_skip": false, + "placeholder": "", + "real_time_refresh": true, + "required": false, + "show": true, + "title_case": false, + "toggle": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "str", + "value": "" + }, + "04_dimension": { + "_input_type": "IntInput", + "advanced": false, + "display_name": "Dimensions", + "dynamic": false, + "info": "Dimensions of the embeddings to generate.", + "list": false, + "list_add_label": "Add More", + "name": "dimension", + "override_skip": false, + "placeholder": "", + "required": false, + "show": true, + "title_case": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "int", + "value": 0 + } + } + } + } + }, + "functionality": "create" + }, + "display_name": "Collection", + "dynamic": false, + "external_options": {}, + "info": "The name of the collection within Astra DB where the vectors will be stored.", + "name": "collection_name", + "options": [], + "options_metadata": [], + "override_skip": false, + "placeholder": "", + "real_time_refresh": true, + "refresh_button": true, + "required": true, + "show": false, + "title_case": false, + "toggle": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "str", + "value": "" + }, + "database_name": { + "_input_type": "DropdownInput", + "advanced": false, + "combobox": true, + "dialog_inputs": { + "fields": { + "data": { + "node": { + "description": "Please allow several minutes for creation to complete.", + "display_name": "Create new database", + "field_order": [ + "01_new_database_name", + "02_cloud_provider", + "03_region" + ], + "name": "create_database", + "template": { + "01_new_database_name": { + "_input_type": "StrInput", + "advanced": false, + "display_name": "Name", + "dynamic": false, + "info": "Name of the new database to create in Astra DB.", + "list": false, + "list_add_label": "Add More", + "load_from_db": false, + "name": "new_database_name", + "override_skip": false, + "placeholder": "", + "required": true, + "show": true, + "title_case": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": false, + "type": "str", + "value": "" + }, + "02_cloud_provider": { + "_input_type": "DropdownInput", + "advanced": false, + "combobox": false, + "dialog_inputs": {}, + "display_name": "Cloud provider", + "dynamic": false, + "external_options": {}, + "info": "Cloud provider for the new database.", + "name": "cloud_provider", + "options": [], + "options_metadata": [], + "override_skip": false, + "placeholder": "", + "real_time_refresh": true, + "required": true, + "show": true, + "title_case": false, + "toggle": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "str", + "value": "" + }, + "03_region": { + "_input_type": "DropdownInput", + "advanced": false, + "combobox": false, + "dialog_inputs": {}, + "display_name": "Region", + "dynamic": false, + "external_options": {}, + "info": "Region for the new database.", + "name": "region", + "options": [], + "options_metadata": [], + "override_skip": false, + "placeholder": "", + "required": true, + "show": true, + "title_case": false, + "toggle": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "str", + "value": "" + } + } + } + } + }, + "functionality": "create" + }, + "display_name": "Database", + "dynamic": false, + "external_options": {}, + "info": "The Database name for the Astra DB instance.", + "name": "database_name", + "options": [], + "options_metadata": [], + "override_skip": false, + "placeholder": "", + "real_time_refresh": true, + "refresh_button": true, + "required": true, + "show": true, + "title_case": false, + "toggle": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "str", + "value": "" + }, + "document": { + "_input_type": "NestedDictInput", + "advanced": false, + "display_name": "Document", + "dynamic": false, + "info": "Single document to insert, e.g. {\"name\": \"Ada\", \"age\": 36}.", + "list": false, + "list_add_label": "Add More", + "name": "document", + "override_skip": false, + "placeholder": "", + "required": false, + "show": false, + "title_case": false, + "tool_mode": true, + "trace_as_input": true, + "trace_as_metadata": true, + "track_in_telemetry": false, + "type": "NestedDict", + "value": {} + }, + "documents": { + "_input_type": "NestedDictInput", + "advanced": false, + "display_name": "Documents", + "dynamic": false, + "info": "List of documents to insert. Provide a JSON list, e.g. [{...}, {...}].", + "list": false, + "list_add_label": "Add More", + "name": "documents", + "override_skip": false, + "placeholder": "", + "required": false, + "show": false, + "title_case": false, + "tool_mode": true, + "trace_as_input": true, + "trace_as_metadata": true, + "track_in_telemetry": false, + "type": "NestedDict", + "value": {} + }, + "environment": { + "_input_type": "DropdownInput", + "advanced": true, + "combobox": true, + "dialog_inputs": {}, + "display_name": "Environment", + "dynamic": false, + "external_options": {}, + "info": "The environment for the Astra DB API Endpoint.", + "name": "environment", + "options": [ + "prod", + "test", + "dev" + ], + "options_metadata": [], + "override_skip": false, + "placeholder": "", + "real_time_refresh": true, + "required": false, + "show": true, + "title_case": false, + "toggle": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "str", + "value": "prod" + }, + "filter_query": { + "_input_type": "NestedDictInput", + "advanced": false, + "display_name": "Filter", + "dynamic": false, + "info": "MongoDB-style filter, e.g. {\"status\": \"active\", \"age\": {\"$gte\": 18}}. Leave empty to match all documents (not recommended for delete/update operations).", + "list": false, + "list_add_label": "Add More", + "name": "filter_query", + "override_skip": false, + "placeholder": "", + "required": false, + "show": true, + "title_case": false, + "tool_mode": true, + "trace_as_input": true, + "trace_as_metadata": true, + "track_in_telemetry": false, + "type": "NestedDict", + "value": {} + }, + "include_similarity": { + "_input_type": "BoolInput", + "advanced": true, + "display_name": "Include Similarity", + "dynamic": false, + "info": "Include the ``$similarity`` score on each returned document (vector searches only).", + "list": false, + "list_add_label": "Add More", + "name": "include_similarity", + "override_skip": false, + "placeholder": "", + "required": false, + "show": true, + "title_case": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "bool", + "value": false + }, + "keyspace": { + "_input_type": "DropdownInput", + "advanced": true, + "combobox": false, + "dialog_inputs": {}, + "display_name": "Keyspace", + "dynamic": false, + "external_options": {}, + "info": "Optional keyspace within Astra DB to use for the collection.", + "name": "keyspace", + "options": [], + "options_metadata": [], + "override_skip": false, + "placeholder": "", + "real_time_refresh": true, + "required": false, + "show": true, + "title_case": false, + "toggle": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "str", + "value": "" + }, + "limit": { + "_input_type": "IntInput", + "advanced": false, + "display_name": "Limit", + "dynamic": false, + "info": "Maximum number of documents to return.", + "list": false, + "list_add_label": "Add More", + "name": "limit", + "override_skip": false, + "placeholder": "", + "required": false, + "show": true, + "title_case": false, + "tool_mode": true, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "int", + "value": 100 + }, + "operation": { + "_input_type": "DropdownInput", + "advanced": false, + "combobox": false, + "dialog_inputs": {}, + "display_name": "Operation", + "dynamic": false, + "external_options": {}, + "info": "Data API operation to run against the selected collection.", + "name": "operation", + "options": [ + "Find", + "Find One", + "Insert One", + "Insert Many", + "Update One", + "Update Many", + "Delete One", + "Delete Many", + "Count Documents", + "Estimated Count" + ], + "options_metadata": [ + { + "icon": "Search" + }, + { + "icon": "SearchCheck" + }, + { + "icon": "Plus" + }, + { + "icon": "CopyPlus" + }, + { + "icon": "Pencil" + }, + { + "icon": "PencilLine" + }, + { + "icon": "Trash" + }, + { + "icon": "Trash2" + }, + { + "icon": "Hash" + }, + { + "icon": "Sigma" + } + ], + "override_skip": false, + "placeholder": "", + "real_time_refresh": true, + "required": false, + "show": true, + "title_case": false, + "toggle": false, + "tool_mode": true, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "str", + "value": "Find" + }, + "ordered": { + "_input_type": "BoolInput", + "advanced": true, + "display_name": "Ordered Insert", + "dynamic": false, + "info": "If enabled, inserts stop at the first error; otherwise errors are collected.", + "list": false, + "list_add_label": "Add More", + "name": "ordered", + "override_skip": false, + "placeholder": "", + "required": false, + "show": false, + "title_case": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "bool", + "value": false + }, + "projection": { + "_input_type": "NestedDictInput", + "advanced": true, + "display_name": "Projection", + "dynamic": false, + "info": "Fields to include (``1``/``true``) or exclude (``0``/``false``), e.g. {\"name\": 1, \"email\": 1}.", + "list": false, + "list_add_label": "Add More", + "name": "projection", + "override_skip": false, + "placeholder": "", + "required": false, + "show": true, + "title_case": false, + "tool_mode": false, + "trace_as_input": true, + "trace_as_metadata": true, + "track_in_telemetry": false, + "type": "NestedDict", + "value": {} + }, + "skip": { + "_input_type": "IntInput", + "advanced": true, + "display_name": "Skip", + "dynamic": false, + "info": "Number of documents to skip.", + "list": false, + "list_add_label": "Add More", + "name": "skip", + "override_skip": false, + "placeholder": "", + "required": false, + "show": true, + "title_case": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "int", + "value": 0 + }, + "sort": { + "_input_type": "NestedDictInput", + "advanced": true, + "display_name": "Sort", + "dynamic": false, + "info": "Sort specification -- ``1`` ascending, ``-1`` descending. Use ``$vector`` or ``$vectorize`` for vector search, e.g. {\"$vectorize\": \"search query\"}.", + "list": false, + "list_add_label": "Add More", + "name": "sort", + "override_skip": false, + "placeholder": "", + "required": false, + "show": true, + "title_case": false, + "tool_mode": false, + "trace_as_input": true, + "trace_as_metadata": true, + "track_in_telemetry": false, + "type": "NestedDict", + "value": {} + }, + "token": { + "_input_type": "SecretStrInput", + "advanced": false, + "display_name": "Astra DB Application Token", + "dynamic": false, + "info": "Authentication token for accessing Astra DB.", + "input_types": [], + "load_from_db": true, + "name": "token", + "override_skip": false, + "password": true, + "placeholder": "", + "real_time_refresh": true, + "required": true, + "show": true, + "title_case": false, + "track_in_telemetry": false, + "type": "str", + "value": "ASTRA_DB_APPLICATION_TOKEN" + }, + "update": { + "_input_type": "NestedDictInput", + "advanced": false, + "display_name": "Update", + "dynamic": false, + "info": "Update operator document, e.g. {\"$set\": {\"status\": \"archived\"}}. Must use update operators (``$set``, ``$inc``, ``$push`` ...).", + "list": false, + "list_add_label": "Add More", + "name": "update", + "override_skip": false, + "placeholder": "", + "required": false, + "show": false, + "title_case": false, + "tool_mode": true, + "trace_as_input": true, + "trace_as_metadata": true, + "track_in_telemetry": false, + "type": "NestedDict", + "value": {} + }, + "upper_bound": { + "_input_type": "IntInput", + "advanced": true, + "display_name": "Count Upper Bound", + "dynamic": false, + "info": "Maximum count astrapy will scan to for ``Count Documents``.", + "list": false, + "list_add_label": "Add More", + "name": "upper_bound", + "override_skip": false, + "placeholder": "", + "required": false, + "show": false, + "title_case": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "int", + "value": 1000 + }, + "upsert": { + "_input_type": "BoolInput", + "advanced": true, + "display_name": "Upsert", + "dynamic": false, + "info": "Insert a new document if no match is found.", + "list": false, + "list_add_label": "Add More", + "name": "upsert", + "override_skip": false, + "placeholder": "", + "required": false, + "show": false, + "title_case": false, + "tool_mode": false, + "trace_as_metadata": true, + "track_in_telemetry": true, + "type": "bool", + "value": false + } + }, + "tool_mode": false + }, "AstraDBGraph": { "base_classes": [ "JSON", @@ -118140,9 +118946,9 @@ ] ], "metadata": { - "num_components": 355, + "num_components": 356, "num_modules": 97 }, - "sha256": "0d5c1c52666c1b5cb69d94ee754844d8c4c94d1d94bfe10b2f97ecf7676f08d0", + "sha256": "1d0ed73df27c720f85d36b6c57ca242031c0f4bae15cd3068adb27ef85e1900b", "version": "0.5.0" } diff --git a/src/lfx/src/lfx/components/datastax/__init__.py b/src/lfx/src/lfx/components/datastax/__init__.py index 5f6f074ccd..b01054811f 100644 --- a/src/lfx/src/lfx/components/datastax/__init__.py +++ b/src/lfx/src/lfx/components/datastax/__init__.py @@ -7,6 +7,7 @@ from lfx.components._importing import import_mod if TYPE_CHECKING: from .astradb_chatmemory import AstraDBChatMemory from .astradb_cql import AstraDBCQLToolComponent + from .astradb_data_api import AstraDBDataAPIComponent from .astradb_graph import AstraDBGraphVectorStoreComponent from .astradb_tool import AstraDBToolComponent from .astradb_vectorize import AstraVectorizeComponent @@ -18,6 +19,7 @@ if TYPE_CHECKING: _dynamic_imports = { "AstraDBCQLToolComponent": "astradb_cql", "AstraDBChatMemory": "astradb_chatmemory", + "AstraDBDataAPIComponent": "astradb_data_api", "AstraDBGraphVectorStoreComponent": "astradb_graph", "AstraDBToolComponent": "astradb_tool", "AstraDBVectorStoreComponent": "astradb_vectorstore", @@ -30,6 +32,7 @@ _dynamic_imports = { __all__ = [ "AstraDBCQLToolComponent", "AstraDBChatMemory", + "AstraDBDataAPIComponent", "AstraDBGraphVectorStoreComponent", "AstraDBToolComponent", "AstraDBVectorStoreComponent", diff --git a/src/lfx/src/lfx/components/datastax/astradb_data_api.py b/src/lfx/src/lfx/components/datastax/astradb_data_api.py new file mode 100644 index 0000000000..5c4c0c81b0 --- /dev/null +++ b/src/lfx/src/lfx/components/datastax/astradb_data_api.py @@ -0,0 +1,614 @@ +"""Astra DB Data API component. + +A thin, modern Langflow component that exposes the full document-based +Data API surface of DataStax Astra DB using **only** the ``astrapy`` SDK. + +The component inherits :class:`lfx.base.datastax.astradb_base.AstraDBBaseComponent` +to reuse the polished database / collection / keyspace selector UI (including +the in-app "Create new database" and "Create new collection" dialogs). +Unlike :class:`AstraDBVectorStoreComponent`, no ``langchain-astradb`` code path +is used at runtime --- every operation goes through ``astrapy`` directly. + +Supported operations (selectable via the operation tab): + +* Find --- ``collection.find`` +* Find One --- ``collection.find_one`` +* Insert One --- ``collection.insert_one`` +* Insert Many --- ``collection.insert_many`` +* Update One --- ``collection.update_one`` +* Update Many --- ``collection.update_many`` +* Delete One --- ``collection.delete_one`` +* Delete Many --- ``collection.delete_many`` +* Count Documents --- ``collection.count_documents`` +* Estimated Count --- ``collection.estimated_document_count`` +""" + +from __future__ import annotations + +from typing import Any + +from astrapy import Collection, DataAPIClient, Database +from astrapy.info import ( + CollectionDefinition, + CollectionVectorOptions, + VectorServiceOptions, +) + +from lfx.base.datastax.astradb_base import AstraDBBaseComponent +from lfx.io import ( + BoolInput, + DropdownInput, + IntInput, + NestedDictInput, + Output, +) +from lfx.log.logger import logger +from lfx.schema.data import Data +from lfx.schema.dataframe import DataFrame + +# Operation option constants -- kept as module-level constants so the UI +# tab values and the dispatcher stay in sync with a single source of truth. +OP_FIND = "Find" +OP_FIND_ONE = "Find One" +OP_INSERT_ONE = "Insert One" +OP_INSERT_MANY = "Insert Many" +OP_UPDATE_ONE = "Update One" +OP_UPDATE_MANY = "Update Many" +OP_DELETE_ONE = "Delete One" +OP_DELETE_MANY = "Delete Many" +OP_COUNT = "Count Documents" +OP_ESTIMATED_COUNT = "Estimated Count" + +ALL_OPERATIONS: tuple[str, ...] = ( + OP_FIND, + OP_FIND_ONE, + OP_INSERT_ONE, + OP_INSERT_MANY, + OP_UPDATE_ONE, + OP_UPDATE_MANY, + OP_DELETE_ONE, + OP_DELETE_MANY, + OP_COUNT, + OP_ESTIMATED_COUNT, +) + +# Per-operation icons surfaced in the operation dropdown. Icon names mirror +# those used elsewhere in Langflow for a consistent look. +OPERATION_ICONS: dict[str, str] = { + OP_FIND: "Search", + OP_FIND_ONE: "SearchCheck", + OP_INSERT_ONE: "Plus", + OP_INSERT_MANY: "CopyPlus", + OP_UPDATE_ONE: "Pencil", + OP_UPDATE_MANY: "PencilLine", + OP_DELETE_ONE: "Trash", + OP_DELETE_MANY: "Trash2", + OP_COUNT: "Hash", + OP_ESTIMATED_COUNT: "Sigma", +} + +# Groups of inputs that each operation needs. Centralising this makes the +# dynamic show/hide logic trivial to audit and extend. +OPERATION_FIELDS: dict[str, set[str]] = { + OP_FIND: {"filter_query", "projection", "sort", "limit", "skip", "include_similarity"}, + OP_FIND_ONE: {"filter_query", "projection", "sort", "include_similarity"}, + OP_INSERT_ONE: {"document"}, + OP_INSERT_MANY: {"documents", "ordered"}, + OP_UPDATE_ONE: {"filter_query", "update", "upsert"}, + OP_UPDATE_MANY: {"filter_query", "update", "upsert"}, + OP_DELETE_ONE: {"filter_query"}, + OP_DELETE_MANY: {"filter_query"}, + OP_COUNT: {"filter_query", "upper_bound"}, + OP_ESTIMATED_COUNT: set(), +} + +# Default count upper bound (``count_documents`` in astrapy requires one). +DEFAULT_COUNT_UPPER_BOUND = 1000 +# Default find limit -- guard against unbounded scans in the UI path. +DEFAULT_FIND_LIMIT = 100 + +# Fields specific to the operation that we toggle ``show`` on. +_OPERATION_TOGGLE_FIELDS: tuple[str, ...] = ( + "filter_query", + "projection", + "sort", + "limit", + "skip", + "include_similarity", + "document", + "documents", + "update", + "upsert", + "ordered", + "upper_bound", +) + + +class AstraDBDataAPIComponent(AstraDBBaseComponent): + """Direct ``astrapy`` Data API access for Astra DB collections. + + Inherits the standard Astra DB selector UI (database / keyspace / + collection, plus the create-new dialogs) from + :class:`AstraDBBaseComponent` and adds a single operation tab that + drives a minimal, operation-specific set of inputs. + """ + + display_name: str = "Astra DB Data API" + description: str = ( + "Run Data API operations (find, filter, insert, update, delete, " + "count) against an Astra DB collection using astrapy directly." + ) + documentation: str = "https://docs.datastax.com/en/astra-db-serverless/api-reference/overview.html" + icon: str = "AstraDB" + name: str = "AstraDBDataAPI" + beta: bool = True + + inputs = [ + *AstraDBBaseComponent.inputs, + DropdownInput( + name="operation", + display_name="Operation", + info="Data API operation to run against the selected collection.", + options=list(ALL_OPERATIONS), + options_metadata=[{"icon": OPERATION_ICONS[op]} for op in ALL_OPERATIONS], + value=OP_FIND, + real_time_refresh=True, + combobox=False, + tool_mode=True, + ), + # -- Query / projection / sort --------------------------------- + NestedDictInput( + name="filter_query", + display_name="Filter", + info=( + "MongoDB-style filter, e.g. " + '{"status": "active", "age": {"$gte": 18}}. ' + "Leave empty to match all documents (not recommended for " + "delete/update operations)." + ), + advanced=False, + tool_mode=True, + ), + NestedDictInput( + name="projection", + display_name="Projection", + info=('Fields to include (``1``/``true``) or exclude (``0``/``false``), e.g. {"name": 1, "email": 1}.'), + advanced=True, + ), + NestedDictInput( + name="sort", + display_name="Sort", + info=( + "Sort specification -- ``1`` ascending, ``-1`` descending. " + "Use ``$vector`` or ``$vectorize`` for vector search, e.g. " + '{"$vectorize": "search query"}.' + ), + advanced=True, + ), + IntInput( + name="limit", + display_name="Limit", + info="Maximum number of documents to return.", + value=DEFAULT_FIND_LIMIT, + advanced=False, + tool_mode=True, + ), + IntInput( + name="skip", + display_name="Skip", + info="Number of documents to skip.", + value=0, + advanced=True, + ), + BoolInput( + name="include_similarity", + display_name="Include Similarity", + info="Include the ``$similarity`` score on each returned document (vector searches only).", + value=False, + advanced=True, + ), + # -- Writes ----------------------------------------------------- + NestedDictInput( + name="document", + display_name="Document", + info='Single document to insert, e.g. {"name": "Ada", "age": 36}.', + show=False, + tool_mode=True, + ), + NestedDictInput( + name="documents", + display_name="Documents", + info="List of documents to insert. Provide a JSON list, e.g. [{...}, {...}].", + show=False, + tool_mode=True, + ), + BoolInput( + name="ordered", + display_name="Ordered Insert", + info="If enabled, inserts stop at the first error; otherwise errors are collected.", + value=False, + show=False, + advanced=True, + ), + NestedDictInput( + name="update", + display_name="Update", + info=( + 'Update operator document, e.g. {"$set": {"status": "archived"}}. ' + "Must use update operators (``$set``, ``$inc``, ``$push`` ...)." + ), + show=False, + tool_mode=True, + ), + BoolInput( + name="upsert", + display_name="Upsert", + info="Insert a new document if no match is found.", + value=False, + show=False, + advanced=True, + ), + # -- Count ------------------------------------------------------ + IntInput( + name="upper_bound", + display_name="Count Upper Bound", + info="Maximum count astrapy will scan to for ``Count Documents``.", + value=DEFAULT_COUNT_UPPER_BOUND, + show=False, + advanced=True, + ), + ] + + outputs = [ + Output(display_name="Data", name="data", method="run_operation"), + Output(display_name="Table", name="dataframe", method="as_dataframe"), + Output(display_name="Raw Result", name="raw", method="raw_result"), + ] + + # ---------------------------------------------------------------------- + # Override collection creation to use astrapy *directly*. + # ---------------------------------------------------------------------- + + @classmethod + async def create_collection_api( # type: ignore[override] + cls, + new_collection_name: str, + token: str, + api_endpoint: str, + environment: str | None = None, + keyspace: str | None = None, + dimension: int | None = None, + embedding_generation_provider: str | None = None, + embedding_generation_model: str | None = None, + ) -> Collection: + """Create a new Astra DB collection using astrapy only. + + The base class implementation in :class:`AstraDBBaseComponent` uses + ``_AstraDBCollectionEnvironment`` from ``langchain-astradb``. This + override instead calls ``astrapy``'s :meth:`Database.create_collection` + directly, building a :class:`CollectionDefinition` when + server-side vectorize options are requested. + + Args: + new_collection_name: Required collection name. + token: Astra DB application token. + api_endpoint: Full Astra DB API endpoint for the target database. + environment: Astra environment (``prod`` / ``test`` / ``dev``). + keyspace: Keyspace to create the collection in. + dimension: Embedding dimension if using "Bring your own" embeddings. + embedding_generation_provider: Display name of a configured + vectorize provider, or ``"Bring your own"``. + embedding_generation_model: Model name for the selected provider. + + Returns: + The newly created :class:`astrapy.Collection`. + + Raises: + ValueError: If ``new_collection_name`` is empty or if an + unsupported provider configuration is supplied. + """ + if not new_collection_name: + msg = "Collection name is required to create a new collection." + raise ValueError(msg) + + env = cls.get_environment(environment) + client = DataAPIClient(environment=env) + database = client.get_database(api_endpoint, token=token, keyspace=keyspace) + + # Build the vector options -- either bring-your-own (dimension only) + # or a server-side vectorize provider. + vector_options: CollectionVectorOptions | None = None + if dimension: + vector_options = CollectionVectorOptions(dimension=dimension) + elif embedding_generation_provider and embedding_generation_provider != "Bring your own": + providers = cls.get_vectorize_providers(token=token, environment=env, api_endpoint=api_endpoint) + provider_key = providers.get(embedding_generation_provider, [None, []])[0] + if provider_key is None: + msg = ( + f"Unknown embedding provider '{embedding_generation_provider}'. " + "Pass a dimension for a bring-your-own collection, or choose a " + "configured vectorize provider." + ) + raise ValueError(msg) + vector_options = CollectionVectorOptions( + service=VectorServiceOptions( + provider=provider_key, + model_name=embedding_generation_model, + ), + ) + + definition = CollectionDefinition(vector=vector_options) if vector_options else None + return database.create_collection( + new_collection_name, + definition=definition, + keyspace=keyspace, + ) + + # ---------------------------------------------------------------------- + # Dynamic UI -- show/hide operation-specific fields. + # ---------------------------------------------------------------------- + + async def update_build_config( + self, + build_config: dict, + field_value: str | dict, + field_name: str | None = None, + ) -> dict: + """Keep the base Astra DB selector behavior and toggle operation fields.""" + build_config = await super().update_build_config(build_config, field_value, field_name) + + # On operation change (or first render), toggle visibility of + # operation-specific inputs. + if field_name == "operation" or field_name is None: + op_value = field_value if field_name == "operation" else build_config.get("operation", {}).get("value") + self._apply_operation_visibility(build_config, op_value or OP_FIND) + + return build_config + + @classmethod + def _apply_operation_visibility(cls, build_config: dict, operation: str) -> None: + """Show/hide operation-specific fields based on the selected operation.""" + active = OPERATION_FIELDS.get(operation, set()) + for field in _OPERATION_TOGGLE_FIELDS: + if field in build_config: + build_config[field]["show"] = field in active + + # ---------------------------------------------------------------------- + # Collection accessor + # ---------------------------------------------------------------------- + + def _get_collection(self) -> Collection: + """Return the resolved astrapy :class:`Collection` for this component.""" + if not self.collection_name: + msg = "No collection selected. Choose or create a collection first." + raise ValueError(msg) + + database: Database = self.get_database_object(api_endpoint=self.get_api_endpoint()) + return database.get_collection(self.collection_name, keyspace=self.get_keyspace()) + + # ---------------------------------------------------------------------- + # Operation dispatch + # ---------------------------------------------------------------------- + + def run_operation(self) -> list[Data]: + """Execute the selected operation and return results as ``list[Data]``. + + For write / count operations where a ``list[Data]`` isn't semantically + meaningful, a single-element list wrapping a summary :class:`Data` + object is returned so the output socket stays consistent. + """ + result = self._dispatch() + self.status = result + return result + + def raw_result(self) -> Data: + """Return the raw operation result as a single :class:`Data` object. + + Useful for inspecting the full ``astrapy`` response envelope + (``inserted_ids``, ``matched_count``, ``modified_count`` ...). + """ + raw = self._dispatch(raw=True) + data = Data(data=raw if isinstance(raw, dict) else {"result": raw}) + self.status = data + return data + + def as_dataframe(self) -> DataFrame: + """Return results as a :class:`DataFrame` for easy tabular display.""" + rows = self.run_operation() + return DataFrame([row.data for row in rows if isinstance(row, Data)]) + + # ---------------------------------------------------------------------- + # Internal dispatcher + # ---------------------------------------------------------------------- + + def _dispatch(self, *, raw: bool = False) -> Any: + operation = getattr(self, "operation", OP_FIND) or OP_FIND + collection = self._get_collection() + + handlers = { + OP_FIND: self._op_find, + OP_FIND_ONE: self._op_find_one, + OP_INSERT_ONE: self._op_insert_one, + OP_INSERT_MANY: self._op_insert_many, + OP_UPDATE_ONE: self._op_update_one, + OP_UPDATE_MANY: self._op_update_many, + OP_DELETE_ONE: self._op_delete_one, + OP_DELETE_MANY: self._op_delete_many, + OP_COUNT: self._op_count, + OP_ESTIMATED_COUNT: self._op_estimated_count, + } + + handler = handlers.get(operation) + if handler is None: + msg = f"Unsupported operation '{operation}'." + raise ValueError(msg) + + try: + return handler(collection, raw=raw) + except Exception as exc: + logger.exception("Astra DB Data API operation failed: %s", operation) + msg = f"Astra DB Data API '{operation}' failed: {exc}" + raise ValueError(msg) from exc + + # -- Handlers ----------------------------------------------------- + + def _op_find(self, collection: Collection, *, raw: bool) -> Any: + find_kwargs = self._build_find_kwargs() + cursor = collection.find(**find_kwargs) + docs = list(cursor) + if raw: + return {"count": len(docs), "documents": docs} + return [Data(data=doc) for doc in docs] + + def _op_find_one(self, collection: Collection, *, raw: bool) -> Any: + kwargs = self._build_find_kwargs(one=True) + doc = collection.find_one(**kwargs) + if raw: + return {"document": doc} + return [Data(data=doc)] if doc else [] + + def _op_insert_one(self, collection: Collection, *, raw: bool) -> Any: + document = self._require_mapping("document", self.document, allow_empty=False) + result = collection.insert_one(document) + payload = {"inserted_id": _stringify(result.inserted_id)} + if raw: + return payload + return [Data(data=payload)] + + def _op_insert_many(self, collection: Collection, *, raw: bool) -> Any: + documents = _coerce_documents(self.documents) + if not documents: + msg = "Insert Many requires a non-empty list of documents." + raise ValueError(msg) + result = collection.insert_many(documents, ordered=bool(self.ordered)) + payload = { + "inserted_ids": [_stringify(_id) for _id in result.inserted_ids], + "inserted_count": len(result.inserted_ids), + } + if raw: + return payload + return [Data(data=payload)] + + def _op_update_one(self, collection: Collection, *, raw: bool) -> Any: + filter_q = self._require_mapping("filter_query", self.filter_query, allow_empty=False) + update_doc = self._require_mapping("update", self.update, allow_empty=False) + result = collection.update_one(filter_q, update_doc, upsert=bool(self.upsert)) + return self._format_update_result(result, raw=raw) + + def _op_update_many(self, collection: Collection, *, raw: bool) -> Any: + filter_q = self._require_mapping("filter_query", self.filter_query, allow_empty=False) + update_doc = self._require_mapping("update", self.update, allow_empty=False) + result = collection.update_many(filter_q, update_doc, upsert=bool(self.upsert)) + return self._format_update_result(result, raw=raw) + + def _op_delete_one(self, collection: Collection, *, raw: bool) -> Any: + filter_q = self._require_mapping("filter_query", self.filter_query, allow_empty=False) + result = collection.delete_one(filter_q) + payload = {"deleted_count": result.deleted_count} + if raw: + return payload + return [Data(data=payload)] + + def _op_delete_many(self, collection: Collection, *, raw: bool) -> Any: + filter_q = self._require_mapping("filter_query", self.filter_query, allow_empty=False) + result = collection.delete_many(filter_q) + payload = {"deleted_count": result.deleted_count} + if raw: + return payload + return [Data(data=payload)] + + def _op_count(self, collection: Collection, *, raw: bool) -> Any: + filter_q = self.filter_query or {} + upper_bound = int(self.upper_bound or DEFAULT_COUNT_UPPER_BOUND) + count = collection.count_documents(filter_q, upper_bound=upper_bound) + payload = {"count": count, "upper_bound": upper_bound} + if raw: + return payload + return [Data(data=payload)] + + def _op_estimated_count(self, collection: Collection, *, raw: bool) -> Any: + count = collection.estimated_document_count() + payload = {"estimated_count": count} + if raw: + return payload + return [Data(data=payload)] + + # -- Helpers ------------------------------------------------------ + + def _build_find_kwargs(self, *, one: bool = False) -> dict[str, Any]: + """Build kwargs for ``find`` / ``find_one``, omitting empty values.""" + kwargs: dict[str, Any] = {"filter": self.filter_query or {}} + + if self.projection: + kwargs["projection"] = self.projection + if self.sort: + kwargs["sort"] = self.sort + if self.include_similarity: + kwargs["include_similarity"] = True + + if not one: + limit = int(self.limit) if self.limit else 0 + if limit > 0: + kwargs["limit"] = limit + skip = int(self.skip) if self.skip else 0 + if skip > 0: + kwargs["skip"] = skip + return kwargs + + @staticmethod + def _format_update_result(result: Any, *, raw: bool) -> Any: + payload = { + "matched_count": getattr(result, "matched_count", None), + "modified_count": getattr(result, "modified_count", None), + "upserted_id": _stringify(getattr(result, "upserted_id", None)), + } + if raw: + return payload + return [Data(data=payload)] + + @staticmethod + def _require_mapping(field_name: str, value: Any, *, allow_empty: bool = True) -> dict: + """Validate that a field contains a mapping; raise a clear error otherwise.""" + if value is None or value == "": + if allow_empty: + return {} + msg = f"Field '{field_name}' is required for this operation." + raise ValueError(msg) + if not isinstance(value, dict): + msg = f"Field '{field_name}' must be a JSON object, got {type(value).__name__}." + raise ValueError(msg) # noqa: TRY004 + if not value and not allow_empty: + msg = f"Field '{field_name}' must not be empty for this operation." + raise ValueError(msg) + return value + + +# ---------------------------------------------------------------------- +# Module-level helpers (kept outside the class for easy unit-testing) +# ---------------------------------------------------------------------- + + +def _stringify(value: Any) -> Any: + """Best-effort stringification of Data API id types for JSON-serialisable output.""" + if value is None: + return None + if isinstance(value, str | int | float | bool): + return value + return str(value) + + +def _coerce_documents(value: Any) -> list[dict[str, Any]]: + """Accept either a ``list[dict]`` or a single ``dict`` and normalise to a list.""" + if value is None or value == "": + return [] + if isinstance(value, dict): + # Singleton -- wrap into a list for convenience. + return [value] + if isinstance(value, list): + bad = [i for i, item in enumerate(value) if not isinstance(item, dict)] + if bad: + msg = f"Documents list contains non-dict entries at index {bad[0]}." + raise ValueError(msg) + return value + msg = f"Documents must be a list of JSON objects, got {type(value).__name__}." + raise ValueError(msg)