Skip to content
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

Add Cohere embeddings integration #3477

Draft
wants to merge 6 commits into
base: master
Choose a base branch
from
Draft

Conversation

morganmcg1
Copy link
Member

@morganmcg1 morganmcg1 commented Jan 23, 2025

This PR adds support for Cohere embeddings integration to Weave.

Changes:

  • Added embeddings wrapper functions for both sync and async clients
  • Added tests for embeddings functionality
  • Generated test cassettes
  • Added proper handling of usage data in meta field

The implementation follows the same pattern as the OpenAI embeddings integration.

Example trace: https://wandb.ai/morgan/cohere-embeddings-example/r/call/01949331-83e2-7e82-b33a-ddd221da25b6

Example usage:

import weave
import cohere

# Initialize weave with a project name
client = weave.init("cohere-embeddings-example")

# Initialize Cohere client
co = cohere.Client("your-api-key")

# Create some example texts
texts = [
    "The quick brown fox jumps over the lazy dog",
    "Machine learning is transforming the world",
    "Artificial intelligence and neural networks"
]

# Get embeddings
response = co.embed(
    texts=texts,
    model="embed-english-v3.0",
    input_type="search_document"
)

# Get the traced calls
calls = list(client.get_calls())

# Example traced output:
# Operation name: weave:///morgan/cohere-embeddings-example/op/cohere.Client.embed:EEbjvTp2FQIfCtuX86kEdYio6gsbggWdimA7g2Xj4VE
# Start time: 2025-01-23 12:33:30.444069+00:00
# End time: 2025-01-23 12:33:30.566521+00:00
# Duration: 0.122452 seconds
# Number of embeddings: 3
# Meta information: {
#   "api_version": {"version": "1"},
#   "billed_units": {"input_tokens": 20.0},
#   "warnings": []
# }

@morganmcg1 morganmcg1 requested a review from a team as a code owner January 23, 2025 12:29
Copy link
Contributor

CLA Assistant Lite bot:
Thank you for your submission, we really appreciate it. Like many open-source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution. You can sign the CLA by just posting a Pull Request Comment same as the below format.


I have read the CLA Document and I hereby sign the CLA


You can retrigger this bot by commenting recheck in this Pull Request

@circle-job-mirror
Copy link

circle-job-mirror bot commented Jan 23, 2025

@morganmcg1 morganmcg1 marked this pull request as draft January 23, 2025 13:22
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.

2 participants