Skip to content

Conversation

@pedrocassalpacheco
Copy link
Contributor

@pedrocassalpacheco pedrocassalpacheco commented Sep 17, 2025

Add ContextualAI Rerank Service Implementation

Overview

This PR adds a complete implementation of ContextualAI rerank service integration to Elasticsearch's inference plugin, following the established patterns used by other inference services (OpenAI, Cohere, etc.).

What's Implemented

  • Complete ContextualAI Service: Full service implementation with rerank task support
  • Service Registration: Integrated with InferencePlugin to register the service
  • Request/Response Handling: Proper HTTP client integration with ContextualAI API
  • Configuration Management: Service settings, task settings, and secret handling
  • Rate Limiting: Integrated with Elasticsearch's rate limiting infrastructure
  • Instruction Support: Added instruction parameter for guiding rerank behavior

Files Added/Modified

  • ContextualAiService.java - Main service implementation
  • ContextualAiRerankModel.java - Model representation
  • ContextualAiRerankServiceSettings.java - Service configuration
  • ContextualAiRerankTaskSettings.java - Task-specific settings with instruction support
  • ContextualAiActionCreator.java - Action creation logic
  • ContextualAiRerankRequest.java - HTTP request handling with debug logging
  • ContextualAiRerankRequestEntity.java - JSON serialization with proper field ordering
  • ContextualAiRerankResponseEntity.java - Response parsing
  • InferencePlugin.java - Service registration

Lastest commit a818ac3

  • This commit has a slightly simpler implementation to support only reranking models.
  • It eliminates some of the complexities of supporting different models, but it works for an initial deployment.

Service registration

{
  "service": "contextualai",
  "service_settings": {
    "model_id": "ctxl-rerank-v2-instruct-multilingual",
    "url": "https://api.contextual.ai/v1/rerank",
    "api_key": "your contextual key"
  }
}

Inference test call

{
  "query": "What is the current enterprise pricing for the RTX 5090 GPU for bulk orders?",
  "input": [
    "Following detailed cost analysis and market research, we have implemented the following changes: AI training clusters will see a 15% uplift in pricing effective Q1 2025.",
    "NVIDIA Enterprise Sales Portal: RTX 5090 GPU bulk pricing for enterprise clients is set at $8,500 per unit as of January 2025",
    "TechAnalytics Research Group: Market analysis indicates stable pricing for RTX 5090 GPUs in Q1 2025, with minor fluctuations expected due to supply chain factors.",
    "I really like popcorn"
  ],
   "task_settings": {
    "instruction": "Prioritize internal sales documents over market analysis reports. More recent documents should be weighted higher. Enterprise portal content supersedes distributor communications."
  }
}

Response

{
    "rerank": [
        {
            "index": 1,
            "relevance_score": 0.9999927
        },
        {
            "index": 2,
            "relevance_score": 0.25158092
        },
        {
            "index": 0,
            "relevance_score": 0.00323565
        },
        {
            "index": 3,
            "relevance_score": 2.4E-7
        }
    ]
}

- Added complete ContextualAI rerank service implementation
- Updated InferencePlugin.java to register ContextualAI service
- Fixed service settings validation for REQUEST vs PERSISTENT contexts
- Updated debug script to clear cluster state on restart

Includes only:
- ContextualAI service files under x-pack/plugin/inference/.../contextualai/
- InferencePlugin.java registration
- debug_elasticsearch.sh modifications
@cla-checker-service
Copy link

cla-checker-service bot commented Sep 17, 2025

💚 CLA has been signed

@elasticsearchmachine elasticsearchmachine added needs:triage Requires assignment of a team area label v9.2.0 external-contributor Pull request authored by a developer outside the Elasticsearch team labels Sep 17, 2025
@davidkyle davidkyle added :ml Machine learning and removed needs:triage Requires assignment of a team area label labels Sep 17, 2025
@davidkyle davidkyle self-assigned this Sep 17, 2025
@elasticsearchmachine elasticsearchmachine added the Team:ML Meta label for the ML team label Sep 17, 2025
@elasticsearchmachine
Copy link
Collaborator

Pinging @elastic/ml-core (Team:ML)

Copy link
Contributor Author

@pedrocassalpacheco pedrocassalpacheco left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@davidkyle
Copy link
Member

@elasticmachine test this please

Copy link
Member

@davidkyle davidkyle left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks great thanks @pedrocassalpacheco

Please add tests for the new classes added under this folder - https://github.com/elastic/elasticsearch/tree/main/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services

You can run the tests from your terminal with this command: ./gradlew :x-pack:plugin:inference:test

@davidkyle
Copy link
Member

@elasticmachine test this please

@pedrocassalpacheco
Copy link
Contributor Author

@davidkyle - thank you for your support. Happy to make any further changes if necessary.

@davidkyle davidkyle changed the title Feature/contextualai clean [ML] Add ContextualAI inference service Sep 25, 2025
@davidkyle
Copy link
Member

@elasticmachine test this please

@davidkyle
Copy link
Member

@elasticmachine test this please

@davidkyle davidkyle enabled auto-merge (squash) September 26, 2025 08:15
Copy link
Member

@davidkyle davidkyle left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@davidkyle davidkyle merged commit 734a6b7 into elastic:main Sep 26, 2025
35 checks passed
szybia added a commit to szybia/elasticsearch that referenced this pull request Sep 26, 2025
…-dls

* upstream/main:
  [ML] Add ContextualAI inference service (elastic#134933)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

>enhancement external-contributor Pull request authored by a developer outside the Elasticsearch team :ml Machine learning Team:ML Meta label for the ML team v9.2.0

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants