🐛 Bug Report: Not getting Completion content for Cohere Reranking

[tkanhe-karini]

Which component is this bug for?

Cohere Instrumentation

📜 Description

When using cohere reranker, in traces I'm getting prompt.0.content but completion.0.content attribute is missing.

👟 Reproduction steps

I'm using the following code:

from traceloop.sdk.decorators import workflow
import cohere

@workflow("rerank")
def rerank(self, query: str, documents: list[str], top_n: int = 1, relevance: float | None = None, model: str = "rerank-english-v3.0") -> list[dict]:
    Traceloop.set_association_properties({"request_id": self.request_id, "service": "rerank"})
    try:
        results = self.client.rerank(query=query, documents=documents, top_n=top_n, model=model, return_documents=True)

        return [
            {"index": res.index, "text": res.document.text, "relevance_score": res.relevance_score}
            for res in results.results
            if res.relevance_score >= relevance
        ]

    except Exception as e:
        logger.exception(f"Error in cohere reranking: {e}")
        raise

👍 Expected behavior

The traces should contain the completion.0.content attributes.

👎 Actual Behavior with Screenshots

Getting the following traces. It contains prompt.0.content but completion.0.content is missing.

[{
        "Timestamp": datetime.datetime(2024, 10, 10, 7, 31, 26, 398631),
        "TraceId": "2640fd7d6055a19b0be9a501e2982cb0",
        "SpanId": "9980dadfcc5096ae",
        "ParentSpanId": "3034be7eb4b5afde",
        "TraceState": "",
        "SpanName": "rerank.workflow",
        "SpanKind": "Internal",
        "ServiceName": "/home/ubuntu/coreml/core_ml_venv/bin/uvicorn",
        "ResourceAttributes": {"service.name": "/home/ubuntu/coreml/core_ml_venv/bin/uvicorn"},
        "ScopeName": "traceloop.tracer",
        "ScopeVersion": "",
        "SpanAttributes": {
            "traceloop.entity.output": "[]",
            "traceloop.workflow.name": "rerank",
            "traceloop.association.properties.question": "what is aws reinvent ?",
            "traceloop.span.kind": "workflow",
            "traceloop.entity.name": "rerank",
            "traceloop.association.properties.request_id": "1728545486771140564",
            "traceloop.association.properties.service": "rerank",
        },
        "Duration": 206660535,
        "StatusCode": "Unset",
        "StatusMessage": "",
        "Events.Timestamp": [],
        "Events.Name": [],
        "Events.Attributes": [],
        "Links.TraceId": [],
        "Links.SpanId": [],
        "Links.TraceState": [],
        "Links.Attributes": [],
    },
    {
        "Timestamp": datetime.datetime(2024, 10, 10, 7, 31, 26, 398806),
        "TraceId": "2640fd7d6055a19b0be9a501e2982cb0",
        "SpanId": "9862a6496dd1189f",
        "ParentSpanId": "9980dadfcc5096ae",
        "TraceState": "",
        "SpanName": "cohere.rerank",
        "SpanKind": "Client",
        "ServiceName": "/home/ubuntu/coreml/core_ml_venv/bin/uvicorn",
        "ResourceAttributes": {"service.name": "/home/ubuntu/coreml/core_ml_venv/bin/uvicorn"},
        "ScopeName": "opentelemetry.instrumentation.cohere",
        "ScopeVersion": "0.32.2",
        "SpanAttributes": {
            "gen_ai.prompt.1.role": "system",
            "llm.request.type": "rerank",
            "traceloop.association.properties.request_id": "1728545486771140564",
            "gen_ai.prompt.2.content": "\n\n“This new capabilty goes beyond the existing self-tuning capabiltes of Amazon Redshift\n‘Senverless, such as machine leaming (ML)-enhanced techniques to adjust your compute,\nmodify the physical schema of the database, create or drop materialized views as needed\n(the one we manage automaticaly, not yours), and vacuum tables. This new capabilty\nbrings more inteligence to decide how to adjust the compute, what background optimizations.\nare required, and when to apply them, and it makes its decisions based on more dimensions\nWe also orchestrate ML-based optimizations for materialized views, lable optimizations, and\n‘workload management when your queries need it\nDuring the preview, you must opt in to enable these Al-driven scaling and optimizations on\nyour workgroups. You configure the system to balance the optimization fo price or\nperformance. There is only one slide to adjust in the console.\nCreate preview workgroup\nWorkgroup\nWorkgro name\nom——\n(© Ti workgroup i bong css in petew Production st ppc. Before cetg hs workup,\neven Ai Serve Tams\n\nPerformance and cost controls i.\n\n© prc perormance rges\n\n© sasecpacty\n\nprice petormnce ets\n\n\\ El so 7 0\nwa ha Trr—\nwn\n",
            "gen_ai.prompt.3.role": "user",
            "gen_ai.prompt.0.content": '\n\nAmazon Redshift adds new Al capabilities, including\nAmazon Q, to boost efficiency and productivity\nNovara 2, 2023\nAWS News Blog\n> 000/000 0 1\nVoiced by Amazon Polly\nAmazon Redshift puts artical inteligence (Al) at your service to optimize efficiencies and\nmake you more productive with two new capabilities that we are launching in preview today.\nFirst, Amazon Redshift Serverless becomes smarter. I scales capacity proactively and\n‘automatically along dimensions such as the complexity of your queries, their frequency, the\nsize of the dataset, and so on to deliver tailored performance optimizations. This allows you\n10 spend less time tuning your data warehouse instances and more time getting value from\nyour data.\nSecond, Amazon Q generative SOL in Amazon Redshift Query Editor generates SQL\nrecommendations from natural language prompts. This helps you to be more productive in\nextracting insights from your data\nLets start with Amazon Redshift Serverloss\nWhen you use Amazon Redshift Serverless, you can now opt in for a preview of Al-driven\nScaling and optimizations. When enabled, the system observes and learns from your usage\npatterns, such as the concurrent number of querles, their complexity, and the time i takes to\nrun them. Then, it automatically optimizes your serverless endpoint to meet your price\nperformance target. Based on AWS intemal testing, this new capabilty may give you upto\nten times better price performance for variable workloads without any manual intervention.\nAldriven scaling and optimizations eliminate the time and effort to manualy resize your\nworkgroup and plan background optimizations based on workload needs. It continually runs\n‘automatic optimizations when they are most valuable for better performance, avoiding\nperformance cfs and time-outs.\n"\n',
            "gen_ai.prompt.2.role": "system",
            "gen_ai.prompt.0.role": "system",
            "gen_ai.prompt.1.content": "\n\nAmazon Q brings generative Al-powered assistance to IT\npros and developers (preview)\nAWS News Blog\n> 000/000 H\nVoiced by Amazon Polly\nToday, we are announcing the preview of Amazon Q. a new type of generative ariicial\ninteligence (Al) powered assistant that is specifically for work and can be tailored to a\ncustomer's business.\nAmazon Q brings a set of capabiltes to support developers and IT professionals. Now you\ncan use Amazon Q to get started building applications on AWS, research best practices\nresolve errors, and get assistance in coding new features for your applications. For example\nAmazon Q Code Transformation can perform Java application upgrades now, rom version 8\nand 1110 version 17.\nAmazon Q i avaiable in multiple areas of AWS to provide quick access to answers and\nideas wherever you work. Here's a quick look at Amazon Q, including in integrated\ndevelopment environment (IDE):\nne\n",
            "gen_ai.prompt.3.content": "what is aws reinvent ?",
            "traceloop.workflow.name": "rerank",
            "traceloop.association.properties.service": "rerank",
            "gen_ai.completion.0.role": "assistant",
            "gen_ai.system": "Cohere",
            "gen_ai.request.model": "rerank-english-v3.0",
        },
        "Duration": 206377616,
        "StatusCode": "Ok",
        "StatusMessage": "",
        "Events.Timestamp": [],
        "Events.Name": [],
        "Events.Attributes": [],
        "Links.TraceId": [],
        "Links.SpanId": [],
        "Links.TraceState": [],
        "Links.Attributes": [],
    }]

🤖 Python Version

3.10

📃 Provide any additional context for the Bug.

dependencies:

cohere==5.10.0
traceloop-sdk==0.32.2

👀 Have you spent some time to check if this bug has been raised before?

• I checked and didn't find similar issue

Are you willing to submit PR?

None

[dosubot]

Hello, @tkanhe-karini! I'm here to help you with any bugs, questions, or contributions you have. Let's work together to solve the issue you're facing.

The completion.0.content attribute is missing in the traces when using the Cohere reranker because the _set_span_rerank_response function does not set this attribute. It sets attributes for each document in the rerank response using the LLM_COMPLETIONS prefix, but not specifically as completion.0.content ^[1].

To continue talking to Dosu, mention @dosu.