[Inference API] Refactor Contextual AI integration and fix multiple issues - #145700

Summary

This pull request refactors the Contextual AI service integration in the inference plugin, fixing critical bugs, removing deprecated settings, and introducing architectural improvements for maintainability and consistency with modern service patterns.

Bugs Fixed

1. Missing NamedWriteableRegistry registration — ContextualAiRerankServiceSettings and ContextualAiRerankTaskSettings were not registered in InferenceNamedWriteablesProvider, preventing models from being serialized/deserialized over the wire.

2. HTTP 503 responses not retried — ContextualAiResponseHandler incorrectly treated HTTP 503 as non-retryable. Now correctly retries 503 alongside 500.

3. Error messages returned HTTP status line instead of body — ContextualAiErrorResponseEntity.fromResponse only returned the status line (e.g., "HTTP/1.1 400 Bad Request"). Replaced with ErrorResponse.fromResponse to read actual response body. ContextualAiErrorResponseEntity deleted entirely.

4. Duplicate topN fallback logic — ContextualAiRerankRequestEntity contained redundant fallback logic that re-read model task settings. Canonical fallback now lives only in ContextualAiRerankRequest.getTopN().

5. Overly restrictive type check in doInfer — doInfer rejected models that were ContextualAiModel but not ContextualAiRerankModel. Broadened to accept any ContextualAiModel.

6. Client-side re-sorting and top-N filtering in response parsing — ContextualAiRerankResponseEntity re-sorted results by relevance score and applied a topN limit client-side. This duplicated server-side behavior. Response parsing now returns results in the order provided by the API without additional filtering.

7. Input type validation missing — validateInputType was a no-op that accepted any input type. Now properly validates using ServiceUtils.validateInputTypeIsUnspecifiedOrInternal.

Functional Changes

1. Removed url service setting — ContextualAiRerankServiceSettings no longer accepts/stores the url field. The endpoint URI is hardcoded as DEFAULT_RERANK_URI (https://api.contextual.ai/v1/rerank) in ContextualAiRerankModel. Transport wire backward-compatibility handled via new TransportVersion contextual_ai_url_service_setting_removed with conditional read/write logic: older nodes still send/receive a URL string, which is read and discarded on deserialization, and the default URL is written for older peers on serialization.

2. Removed hardcoded default instruction — getInstruction() previously always returned "Rerank the given documents based on their relevance to the query." when unset. Now returns null, and the instruction field is omitted from API requests when unset.

3. Removed document field from response parsing — RankedDocEntry no longer parses the document field from the API response. RankedDoc objects always pass null for the document text, since return_documents was never sent to the ContextualAI API.

4. Added updateServiceSettings implementation — ContextualAiRerankServiceSettings.updateServiceSettings now supports updating mutable fields (currently only rate_limit). Immutable fields like model_id are preserved from the original settings.

5. Changed doInfer visibility — doInfer changed from public to protected, consistent with the SenderService base class contract.

Refactoring

1. Introduced ContextualAiServiceSettings abstract class — Replaces ContextualAiRateLimitServiceSettings interface, encapsulating common model ID and rate-limit settings with shared writeTo/StreamInput, toXContent, fromMap, and updateCommonSettings logic via a CommonSettings record.

2. Introduced ContextualAiUtils — Centralizes shared transport-version constants (ML_INFERENCE_CONTEXTUAL_AI_ADDED, ML_INFERENCE_CONTEXTUAL_AI_URL_SERVICE_SETTING_REMOVED).

3. Simplified ContextualAiModel — Removed direct apiKey and rateLimitServiceSettings fields; now delegates to getServiceSettings() and getSecretSettings(). Stores the request URI directly. Added covariant getServiceSettings() / getSecretSettings() overrides returning ContextualAiServiceSettings and DefaultSecretSettings respectively.

4. Decoupled ContextualAiRerankRequestEntity from ContextualAiRerankModel — Now takes a plain String modelId instead of the model object. Field constants changed from private to public for test assertions.

5. Simplified ContextualAiRerankRequest — Removed separate instruction constructor parameter; instruction is now retrieved from model.getTaskSettings().getInstruction() at request-creation time. Removed decorateWithAuth helper method and debug logging with try/catch wrapping. Inlined JSON serialization directly.

6. Simplified ContextualAiRerankResponseEntity — Removed ContextualAiRerankRequest parameter from fromResponse; response parsing no longer depends on the request object.

7. Inlined createModel helper — Deleted private helper in ContextualAiService.parseRequestConfig, using retrieveModelCreatorFromMapOrThrow(...).createFromMaps(...) directly for consistency with other services.

8. Added covariant return types — ContextualAiRerankModel.getTaskSettings() returns ContextualAiRerankTaskSettings and getServiceSettings() returns ContextualAiRerankServiceSettings, eliminating scattered casts in callers.

9. Optimized ContextualAiRerankModel.of() — Returns the same model instance when request task settings are empty or produce no effective change, avoiding unnecessary object creation.

10. Optimized ContextualAiRerankTaskSettings.of() — Returns the original settings instance when the merged result equals the original, avoiding unnecessary object creation.

Tests

1. New ContextualAiActionCreatorTests — Integration-style tests using MockWebServer verifying end-to-end request construction (headers, body fields, auth), task settings overrides, request-level topN priority, and error handling for invalid response formats.

2. New ContextualAiRerankRequestEntityTests — Unit tests verifying XContent serialization with all fields, required-only fields, and null-guard behavior on required parameters.

3. New ContextualAiRerankRequestTests — Unit tests verifying HTTP request construction (method, URI, headers, auth, body), task settings propagation, and request-level topN override behavior.

4. New ContextualAiRerankModelTests — Unit tests verifying model construction from maps, ContextualAiRerankModel.of() merge semantics including identity-return optimizations.

5. New ContextualAiRerankServiceSettingsTests — BWC wire serialization tests, fromMap parsing (required/optional fields, validation errors), updateServiceSettings behavior, and XContent round-tripping.

6. New ContextualAiRerankTaskSettingsTests — BWC wire serialization tests, fromMap parsing with validation (invalid types, zero/negative topN), of() merge semantics, updatedTaskSettings, isEmpty, and XContent round-tripping.

7. New ContextualAiRerankResponseEntityTests — Unit tests verifying response parsing for single/multiple items, empty results, floating-point precision, order preservation, unknown fields tolerance, and error cases (missing required fields, malformed JSON).

8. Updated ContextualAiResponseHandlerTests — Improved tests to verify error body content is included in exception messages, added 503 retryable assertion, extracted shared test fixtures.

9. New ContextualAiRerankTestFixtures — Shared test constants for inference entity ID, model ID, API key, documents, query, rate limits, and parameterized test values used across all test classes.

Testing

RQ
{
    "model": "ctxl-rerank-v2-instruct-multilingual-mini",
    "query": "star wars main character",
    "documents": [
        "luke",
        "like",
        "leia",
        "chewy",
        "r2d2",
        "star",
        "wars"
    ],
    "return_documents": true
}
RS
{
    "detail": [
        {
            "msg": "Extra inputs are not permitted",
            "loc": [
                "return_documents"
            ]
        }
    ]
}

RQ
{
    "service": "contextualai",
    "service_settings": {
        "url": "some_url",
        "api_key": "{{contextual-ai-api-key}}",
        "model_id": "ctxl-rerank-v2-instruct-multilingual-mini"
    },
    "task_settings": {
        "instruction": "sort by relevancy",
        "top_n": 2
    }
}
RS
{
    "error": {
        "root_cause": [
            {
                "type": "status_exception",
                "reason": "Configuration contains settings [{url=some_url}] unknown to the [contextualai] service"
            }
        ],
        "type": "status_exception",
        "reason": "Configuration contains settings [{url=some_url}] unknown to the [contextualai] service"
    },
    "status": 400
}

RQ
{
    "service": "contextualai",
    "service_settings": {
        "api_key": "{{contextual-ai-api-key}}",
        "model_id": "ctxl-rerank-v2-instruct-multilingual-mini"
    },
    "task_settings": {
        "instruction": "sort by relevancy",
        "top_n": 2
    }
}
RS
{
    "inference_id": "contextual-ai-rerank",
    "task_type": "rerank",
    "service": "contextualai",
    "service_settings": {
        "model_id": "ctxl-rerank-v2-instruct-multilingual-mini",
        "rate_limit": {
            "requests_per_minute": 1000
        }
    },
    "task_settings": {
        "top_n": 2,
        "instruction": "sort by relevancy"
    }
}

RQ
{
    "input": [
        "luke",
        "like",
        "leia",
        "chewy",
        "r2d2",
        "star",
        "wars"
    ],
    "query": "star wars main character",
    "top_n": 2,
    "task_settings": {
        "instruction": "classic trilogy only",
        "top_n": 4
    }
}
RS
{
    "rerank": [
        {
            "index": 4,
            "relevance_score": 0.9320585
        },
        {
            "index": 2,
            "relevance_score": 0.9157031
        }
    ]
}

RQ
{
    "service": "contextualai",
    "service_settings": {
        "api_key": "{{contextual-ai-api-key}}",
        "model_id": "ctxl-rerank-v2-instruct-multilingual-mini"
    }
}
RS
{
    "inference_id": "contextual-ai-rerank-2",
    "task_type": "rerank",
    "service": "contextualai",
    "service_settings": {
        "model_id": "ctxl-rerank-v2-instruct-multilingual-mini",
        "rate_limit": {
            "requests_per_minute": 1000
        }
    }
}

RQ
{
    "input": [
        "luke",
        "like",
        "leia",
        "chewy",
        "r2d2",
        "star",
        "wars"
    ],
    "query": "star wars main character"
}
RS
{
    "rerank": [
        {
            "index": 4,
            "relevance_score": 0.9514682
        },
        {
            "index": 2,
            "relevance_score": 0.9389061
        },
        {
            "index": 0,
            "relevance_score": 0.9089205
        },
        {
            "index": 5,
            "relevance_score": 0.73254484
        },
        {
            "index": 1,
            "relevance_score": 0.7259128
        },
        {
            "index": 3,
            "relevance_score": 0.6090389
        },
        {
            "index": 6,
            "relevance_score": 0.5869658
        }
    ]
}

References

• Original PR: [ML] Add ContextualAI inference service #134933
• Closes elastic/search-team#13715
• Closes https://github.com/elastic/search-team/issues/13766