Cohere Java Integration with Future AGI Tracing

Prerequisites

Complete the Java SDK setup first.

Installation

<dependency>
    <groupId>com.github.future-agi.traceAI</groupId>
    <artifactId>traceai-java-cohere</artifactId>
    <version>main-SNAPSHOT</version>
</dependency>

implementation 'com.github.future-agi.traceAI:traceai-java-cohere:main-SNAPSHOT'

You also need the Cohere Java SDK:

<dependency>
    <groupId>com.cohere</groupId>
    <artifactId>cohere-java</artifactId>
    <version>1.5.0</version>
</dependency>

implementation 'com.cohere:cohere-java:1.5.0'

Wrap the client

import ai.traceai.TraceAI;
import ai.traceai.cohere.TracedCohereClient;
import com.cohere.api.Cohere;

TraceAI.initFromEnvironment();

Cohere client = Cohere.builder()
    .token(System.getenv("COHERE_API_KEY"))
    .build();

TracedCohereClient traced = new TracedCohereClient(client);

Chat

import com.cohere.api.requests.ChatRequest;
import com.cohere.api.types.NonStreamedChatResponse;

NonStreamedChatResponse response = traced.chat(ChatRequest.builder()
    .message("What is the capital of France?")
    .model("command-r-plus")
    .temperature(0.7)
    .build());

System.out.println(response.getText());

Span created: “Cohere Chat” with kind LLM

Embeddings

import com.cohere.api.requests.EmbedRequest;
import com.cohere.api.types.EmbedResponse;

EmbedResponse response = traced.embed(EmbedRequest.builder()
    .texts(List.of("Hello world", "Goodbye world"))
    .model("embed-english-v3.0")
    .inputType(EmbedInputType.SEARCH_DOCUMENT)
    .build());

// EmbedResponse is a union type - use the visitor pattern to access results
response.visit(new EmbedResponse.Visitor<Void>() {
    @Override
    public Void visitEmbeddingsFloats(EmbedFloatsResponse floats) {
        System.out.println("Vectors: " + floats.getEmbeddings().size());
        return null;
    }

    @Override
    public Void visitEmbeddingsByType(EmbedByTypeResponse byType) {
        System.out.println("Vectors: " + byType.getEmbeddings().getFloat_().size());
        return null;
    }

    @Override
    public Void _visitUnknown(Object unknown) {
        return null;
    }
});

Span created: “Cohere Embed” with kind EMBEDDING

Reranking

Cohere is the only Java integration with reranking. Uses FISpanKind.RERANKER.

import com.cohere.api.requests.RerankRequest;
import com.cohere.api.types.RerankResponse;

RerankResponse response = traced.rerank(RerankRequest.builder()
    .query("What is the capital of France?")
    .documents(List.of(
        RerankRequestDocumentsItem.of("Paris is the capital of France."),
        RerankRequestDocumentsItem.of("Berlin is the capital of Germany."),
        RerankRequestDocumentsItem.of("The Eiffel Tower is in Paris.")
    ))
    .model("rerank-english-v3.0")
    .topN(2)
    .build());

for (var result : response.getResults()) {
    System.out.println("Index: " + result.getIndex() + ", Score: " + result.getRelevanceScore());
}

Span created: “Cohere Rerank” with kind RERANKER

What gets captured

Chat spans

Embedding spans

Reranker spans

Accessing the original client

Cohere original = traced.unwrap();