Skip to content

feat: add Cohere embedding integration #1305

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 8 commits into
base: develop
Choose a base branch
from

Conversation

bwook00
Copy link
Contributor

@bwook00 bwook00 commented Jul 24, 2025

Description

Add Cohere Embedding provider

@Pouyanpi,
I added it along with #1304.
(It probably conflicts with 1304)

Checklist

  • I've read the CONTRIBUTING guidelines.
  • I've updated the documentation if applicable.
  • I've added tests if applicable.
  • @mentions of the person or team responsible for reviewing proposed changes.

Copy link

Documentation preview

https://nvidia.github.io/NeMo-Guardrails/review/pr-1305

@Pouyanpi Pouyanpi requested review from Copilot and Pouyanpi July 29, 2025 12:00
Copy link

@Copilot Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull Request Overview

This PR adds Cohere as a new embedding provider to the NeMo Guardrails framework, enabling users to use Cohere's embedding models for document encoding tasks.

  • Implements CohereEmbeddingModel class with both synchronous and asynchronous encoding methods
  • Adds comprehensive test coverage for the new Cohere integration
  • Updates documentation to include Cohere in the supported embedding providers table

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
nemoguardrails/embeddings/providers/cohere.py Implements the CohereEmbeddingModel class with sync/async encoding capabilities
nemoguardrails/embeddings/providers/init.py Registers the new Cohere embedding provider in the system
tests/test_embeddings_cohere.py Adds comprehensive test coverage for Cohere embedding functionality
tests/test_configs/with_cohere_embeddings/config.yml Test configuration file for Cohere embeddings setup
tests/test_configs/with_cohere_embeddings/config.co Test flow configuration for validating Cohere integration
docs/user-guides/configuration-guide.md Updates documentation to include Cohere in supported providers table
Comments suppressed due to low confidence (1)

tests/test_embeddings_cohere.py:70

  • Function name test_live_query is duplicated. This function should have a distinct name like test_live_query_sync to differentiate it from the async version on line 52.
def test_live_query(app):

# async_client_var.set(async_client)
#
# # Make embedding request to Cohere API
# embeddings = await async_client.embed(texts=documents, model=self.model, input_type=self.input_type).embeddings
Copy link
Preview

Copilot AI Jul 29, 2025

Choose a reason for hiding this comment

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

The commented-out async implementation should be removed or properly documented with a TODO/FIXME comment explaining when it might be re-enabled, rather than leaving dead code in the implementation.

Copilot uses AI. Check for mistakes.

@Pouyanpi
Copy link
Collaborator

@bwook00 Thank you for your PR 🎉

Would you please enable pre-commits and apply it? Please see contributing guidelines.

once it is enabled you can do:

poetry run pre-commit run --all-files

@codecov-commenter
Copy link

Codecov Report

❌ Patch coverage is 42.85714% with 16 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.40%. Comparing base (bee719b) to head (025c4f7).

Files with missing lines Patch % Lines
nemoguardrails/embeddings/providers/cohere.py 38.46% 16 Missing ⚠️
Additional details and impacted files
@@             Coverage Diff             @@
##           develop    #1305      +/-   ##
===========================================
- Coverage    70.45%   70.40%   -0.05%     
===========================================
  Files          161      162       +1     
  Lines        16214    16241      +27     
===========================================
+ Hits         11423    11434      +11     
- Misses        4791     4807      +16     
Flag Coverage Δ
python 70.40% <42.85%> (-0.05%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
nemoguardrails/embeddings/providers/__init__.py 96.42% <100.00%> (+0.13%) ⬆️
nemoguardrails/embeddings/providers/cohere.py 38.46% <38.46%> (ø)
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@bwook00
Copy link
Contributor Author

bwook00 commented Jul 30, 2025

@Pouyanpi Thank you for your kind review 😃
Along with this branch, I applied pre-commit too #1304

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants