diff --git a/docsgpt/api/connector/connections.py b/docsgpt/api/connector/connections.py index e3e30ce8..1aa03e8e 100644 --- a/docsgpt/api/connector/connections.py +++ b/docsgpt/api/connector/connections.py @@ -230,8 +230,8 @@ class ConnectionSetup(Resource): @api.doc( description=( "Apply the connect wizard's choices: {create_tools, allow_writes?, tool_permissions?, " - "sync?: {items, frequency, name?}}. allow_writes points GitHub's tool at its write endpoint. " - "Honours an Idempotency-Key header for the sync." + "sync?: {items, frequency, name?, config?}}. allow_writes points GitHub's tool at its write " + "endpoint. config is the synced source's retrieval settings, validated like an upload's. Honours an Idempotency-Key header for the sync." ) ) def post(self, connection_id: str): @@ -298,11 +298,19 @@ class ConnectionSetup(Resource): def _start_sync(user_id: str, row: dict, sync: dict): """Queue the first ingest of a source synced from ``row``. - Returns the source summary, or ``(message, status)`` on a bad request. + Args: + user_id: The connection's owner. + row: The connection row. + sync: ``{items, frequency, name?, config?}``. ``config`` is the source's + retrieval settings (a ``SourceConfig``), applied as an upload's are. + + Returns: + The source summary, or ``(message, status)`` on a bad request. """ from docsgpt.api.user.sources.upload import ( _claim_task_or_get_cached, _derive_source_id, + _parse_source_config, _read_idempotency_key, _scoped_idempotency_key, ) @@ -318,6 +326,9 @@ def _start_sync(user_id: str, row: dict, sync: dict): if frequency not in _FREQUENCIES: return ("Unknown sync frequency", 400) name = (sync.get("name") or "").strip() + source_config, config_error = _parse_source_config(sync.get("config")) + if config_error is not None: + return ("Invalid source config", 400) if definition.sync_ingestor == "github": from docsgpt.parser.remote.github_loader import GitHubLoader @@ -352,6 +363,7 @@ def _start_sync(user_id: str, row: dict, sync: dict): "folder_ids": folder_ids, "recursive": bool(items.get("recursive", True)), "sync_frequency": frequency, + "config": source_config, } else: if definition.sync_ingestor == "linear": @@ -371,6 +383,7 @@ def _start_sync(user_id: str, row: dict, sync: dict): "loader": definition.sync_ingestor, "connection_id": str(row["id"]), "sync_frequency": frequency, + "config": source_config, } idempotency_key, _ = _read_idempotency_key() scoped_key = _scoped_idempotency_key(idempotency_key, user_id) diff --git a/tests/connectors/test_connections_api.py b/tests/connectors/test_connections_api.py index b97428d0..b81f8c06 100644 --- a/tests/connectors/test_connections_api.py +++ b/tests/connectors/test_connections_api.py @@ -144,6 +144,79 @@ class TestSetupSync: # Drive has no tools; nothing was created. assert resp.get_json()["tools"] == [] + def test_retrieval_settings_apply_to_the_synced_source(self, app, pg_conn): + from docsgpt.api.connector.connections import ConnectionSetup + + cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={"token_info": {}}) + config = {"chunking": {"strategy": "markdown", "max_tokens": 800}, "retrieval": {"chunks": 4}} + with _db(pg_conn), patch("docsgpt.api.user.tasks.ingest_connector_task.apply_async", + return_value=MagicMock(id="t")) as apply: + resp = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup", body={ + "sync": {"items": {"folder_ids": ["f1"]}, "config": config}, + }, args=[cid]) + assert resp.status_code == 200 + sent = apply.call_args.kwargs["kwargs"]["config"] + # Validated and filled in like an upload's config. + assert sent["chunking"]["strategy"] == "markdown" + assert sent["chunking"]["max_tokens"] == 800 + assert sent["retrieval"]["chunks"] == 4 + assert sent["kind"] == "classic" + + def test_remote_sync_carries_retrieval_settings(self, app, pg_conn): + from docsgpt.api.connector.connections import ConnectionSetup + + cid = _connection(pg_conn, provider="s3", secrets={"credentials": { + "aws_access_key_id": "AKIA", "aws_secret_access_key": "shh-secret"}}) + with _db(pg_conn), patch("docsgpt.api.user.tasks.ingest_remote.apply_async", + return_value=MagicMock(id="t")) as apply: + resp = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup", body={ + "sync": {"items": {"bucket": "docs"}, "config": {"retrieval": {"chunks": 3}}}, + }, args=[cid]) + assert resp.status_code == 200 + assert apply.call_args.kwargs["kwargs"]["config"]["retrieval"]["chunks"] == 3 + + def test_no_retrieval_settings_keeps_the_defaults(self, app, pg_conn): + from docsgpt.api.connector.connections import ConnectionSetup + + cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={"token_info": {}}) + with _db(pg_conn), patch("docsgpt.api.user.tasks.ingest_connector_task.apply_async", + return_value=MagicMock(id="t")) as apply: + resp = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup", body={ + "sync": {"items": {"folder_ids": ["f1"]}}, + }, args=[cid]) + assert resp.status_code == 200 + assert apply.call_args.kwargs["kwargs"]["config"] is None + + @pytest.mark.parametrize("config", [ + {"retrieval": {"chunks": 0}}, + {"unknown": True}, + {"retrieval": {"chunks": 6, "prescreen": {"candidate_k": 4, "max_keep": 8}}}, + "not an object", + ]) + def test_invalid_retrieval_settings_are_refused_before_queueing(self, app, pg_conn, config): + from docsgpt.api.connector.connections import ConnectionSetup + + @contextmanager + def _yield(): + yield pg_conn + + cid = _connection(pg_conn, provider="google_drive", auth_kind="oauth", secrets={"token_info": {}}) + headers = {"Idempotency-Key": f"setup-config-{hash(str(config))}"} + with _db(pg_conn), patch("docsgpt.api.user.sources.upload.db_session", _yield), patch( + "docsgpt.api.user.tasks.ingest_connector_task.apply_async", return_value=MagicMock(id="t"), + ) as apply: + bad = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup", body={ + "sync": {"items": {"folder_ids": ["f1"]}, "config": config}, + }, headers=headers, args=[cid]) + apply.assert_not_called() + # The refusal left the key free for the corrected retry. + fixed = _call(app, ConnectionSetup, "post", f"/api/connections/{cid}/setup", body={ + "sync": {"items": {"folder_ids": ["f1"]}}, + }, headers=headers, args=[cid]) + assert bad.status_code == 400 + assert bad.get_json()["error"] == "Invalid source config" + assert fixed.status_code == 200 + def test_nothing_picked_is_a_bad_request(self, app, pg_conn): from docsgpt.api.connector.connections import ConnectionSetup