Cohere Java Integration with Future AGI Tracing

Prerequisites

Complete the Java SDK setup first. Cohere preambles, reranking queries and tool-call details can appear in span attributes even when input or output hiding is enabled. Error messages and exception events are also recorded. Avoid tracing sensitive content unless you have separately addressed these fields.

Installation

Add traceai-java-cohere using the same traceAI release and repository as the core dependency in the Java SDK setup.

The wrapper declares Cohere Java SDK 1.9.3 as a provided dependency, so add it to your application:

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

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

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.7f)
    .build());

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

Span created: “Cohere Chat” with kind LLM. Open your project in Future AGI to inspect the span.

Embeddings

import com.cohere.api.requests.EmbedRequest;
import com.cohere.api.types.EmbedResponse;
import com.cohere.api.types.EmbedInputType;
import com.cohere.api.types.EmbedFloatsResponse;
import com.cohere.api.types.EmbedByTypeResponse;
import java.util.List;

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) {
        byType.getEmbeddings().getFloat().ifPresent(vectors ->
            System.out.println("Vectors: " + vectors.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;
import com.cohere.api.types.RerankRequestDocumentsItem;
import java.util.List;

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();