langtrace-python-sdk 2.0.1__py3-none-any.whl → 2.0.2__py3-none-any.whl

langtrace_python_sdk/instrumentation/cohere/patch.py

@@ -23,7 +23,9 @@ from opentelemetry.trace.status import Status, StatusCode
 
 from langtrace_python_sdk.constants.instrumentation.cohere import APIS
 from langtrace_python_sdk.constants.instrumentation.common import (
-    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY, SERVICE_PROVIDERS)
+    LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
+    SERVICE_PROVIDERS,
+)
 
 
 def rerank(original_method, version, tracer):
@@ -103,9 +105,7 @@ def rerank(original_method, version, tracer):
                                 else 0
                             ),
                         }
-                        span.set_attribute(
-                            "llm.token.counts", json.dumps(usage_dict)
-                        )
+                        span.set_attribute("llm.token.counts", json.dumps(usage_dict))
 
             span.set_status(StatusCode.OK)
             span.end()
@@ -187,9 +187,7 @@ def embed(original_method, version, tracer):
                                 else 0
                             ),
                         }
-                        span.set_attribute(
-                            "llm.token.counts", json.dumps(usage_dict)
-                        )
+                        span.set_attribute("llm.token.counts", json.dumps(usage_dict))
 
             span.set_status(StatusCode.OK)
             span.end()
@@ -226,9 +224,7 @@ def chat_create(original_method, version, tracer):
                         item.get("role") if item.get("role") is not None else "USER"
                     ),
                     "content": (
-                        item.get("message")
-                        if item.get("message") is not None
-                        else ""
+                        item.get("message") if item.get("message") is not None else ""
                     ),
                 }
                 for item in chat_history
@@ -312,7 +308,11 @@ def chat_create(original_method, version, tracer):
                 span.set_attribute("llm.is_search_required", result.is_search_required)
 
             if kwargs.get("stream") is False or kwargs.get("stream") is None:
-                if hasattr(result, "text") and result.text is not None and result.text != "":
+                if (
+                    hasattr(result, "text")
+                    and result.text is not None
+                    and result.text != ""
+                ):
                     if (
                         hasattr(result, "chat_history")
                         and result.chat_history is not None
@@ -321,8 +321,7 @@ def chat_create(original_method, version, tracer):
                             {
                                 "role": (
                                     item.role
-                                    if hasattr(item, "role")
-                                    and item.role is not None
+                                    if hasattr(item, "role") and item.role is not None
                                     else "USER"
                                 ),
                                 "content": (
@@ -336,9 +335,7 @@ def chat_create(original_method, version, tracer):
                         ]
                         span.set_attribute("llm.responses", json.dumps(responses))
                     else:
-                        responses = [
-                            {"role": "CHATBOT", "content": result.text}
-                        ]
+                        responses = [{"role": "CHATBOT", "content": result.text}]
                         span.set_attribute("llm.responses", json.dumps(responses))
                 elif hasattr(result, "tool_calls") and result.tool_calls is not None:
                     tool_calls = []
@@ -422,9 +419,7 @@ def chat_stream(original_method, version, tracer):
                         item.get("role") if item.get("role") is not None else "USER"
                     ),
                     "content": (
-                        item.get("message")
-                        if item.get("message") is not None
-                        else ""
+                        item.get("message") if item.get("message") is not None else ""
                     ),
                 }
                 for item in chat_history
@@ -485,7 +480,7 @@ def chat_stream(original_method, version, tracer):
             # stringify the list of objects
             attributes.llm_tool_results = json.dumps(kwargs.get("tool_results"))
 
-        span = tracer.start_span(APIS["CHAT_CREATE"]["METHOD"], kind=SpanKind.CLIENT)
+        span = tracer.start_span(APIS["CHAT_STREAM"]["METHOD"], kind=SpanKind.CLIENT)
         for field, value in attributes.model_dump(by_alias=True).items():
             if value is not None:
                 span.set_attribute(field, value)
@@ -503,19 +498,28 @@ def chat_stream(original_method, version, tracer):
                         Event.STREAM_OUTPUT.value, {"response": "".join(content)}
                     )
 
-                    if hasattr(event, "finish_reason") and event.finish_reason == "COMPLETE":
+                    if (
+                        hasattr(event, "finish_reason")
+                        and event.finish_reason == "COMPLETE"
+                    ):
                         response = event.response
 
                         if (hasattr(response, "generation_id")) and (
                             response.generation_id is not None
                         ):
-                            span.set_attribute("llm.generation_id", response.generation_id)
-                        if (hasattr(response, "response_id")) and (response.response_id is not None):
+                            span.set_attribute(
+                                "llm.generation_id", response.generation_id
+                            )
+                        if (hasattr(response, "response_id")) and (
+                            response.response_id is not None
+                        ):
                             span.set_attribute("llm.response_id", response.response_id)
                         if (hasattr(response, "is_search_required")) and (
                             response.is_search_required is not None
                         ):
-                            span.set_attribute("llm.is_search_required", response.is_search_required)
+                            span.set_attribute(
+                                "llm.is_search_required", response.is_search_required
+                            )
 
                         # Set the response attributes
                         if hasattr(response, "text") and response.text is not None:
@@ -540,12 +544,16 @@ def chat_stream(original_method, version, tracer):
                                     }
                                     for item in response.chat_history
                                 ]
-                                span.set_attribute("llm.responses", json.dumps(responses))
+                                span.set_attribute(
+                                    "llm.responses", json.dumps(responses)
+                                )
                             else:
                                 responses = [
                                     {"role": "CHATBOT", "content": response.text}
                                 ]
-                                span.set_attribute("llm.responses", json.dumps(responses))
+                                span.set_attribute(
+                                    "llm.responses", json.dumps(responses)
+                                )
 
                         # Get the usage
                         if hasattr(response, "meta") and response.meta is not None:

langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "2.0.1"
+__version__ = "2.0.2"

{langtrace_python_sdk-2.0.1.dist-info → langtrace_python_sdk-2.0.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: langtrace-python-sdk
-Version: 2.0.1
+Version: 2.0.2
 Summary: Python SDK for LangTrace
 Project-URL: Homepage, https://github.com/Scale3-Labs/langtrace-python-sdk
 Author-email: Scale3 Labs <engineering@scale3labs.com>

{langtrace_python_sdk-2.0.1.dist-info → langtrace_python_sdk-2.0.2.dist-info}/RECORD

@@ -37,7 +37,7 @@ examples/qdrant_example/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG
 examples/qdrant_example/basic.py,sha256=2JDMsBllnRvVDuuZ91FkAkpXgnavUP4mqXa51CtfHP8,1421
 langtrace_python_sdk/__init__.py,sha256=q2TAsxGIv4wFIQA7G-HGxlnzO89bhPU5TgI42sPHLj0,740
 langtrace_python_sdk/langtrace.py,sha256=HuGbzXM_oqIzpfFDyV2nxuEe4C3SCxDlS1qFzswXzY4,4762
-langtrace_python_sdk/version.py,sha256=wAxkK8w13vqoF47A8iqWdSlIgRRXmZiQ0R4wePZfzhs,22
+langtrace_python_sdk/version.py,sha256=tATvJM5shAzfspHYjdVwpV2w3-gDA119NlEYi5X2lFY,22
 langtrace_python_sdk/constants/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langtrace_python_sdk/constants/exporter/langtrace_exporter.py,sha256=5MNjnAOg-4am78J3gVMH6FSwq5N8TOj72ugkhsw4vi0,46
 langtrace_python_sdk/constants/instrumentation/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -60,7 +60,7 @@ langtrace_python_sdk/instrumentation/chroma/instrumentation.py,sha256=nT6PS6bsrI
 langtrace_python_sdk/instrumentation/chroma/patch.py,sha256=gGBX-zLQcOWUWt596PudOJftZcHge0mA_X1ltW3C800,2802
 langtrace_python_sdk/instrumentation/cohere/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langtrace_python_sdk/instrumentation/cohere/instrumentation.py,sha256=XAs95fPBh2TTd6E1qUbi-FUg7PClW3Y6P1eOq2f38VA,2228
-langtrace_python_sdk/instrumentation/cohere/patch.py,sha256=US6i0Gn4LEEV-Au87q0Scje9JOIiO3H84kAJCEhb4kU,26381
+langtrace_python_sdk/instrumentation/cohere/patch.py,sha256=hdSWArP0cdfGs0Q5KhkaowkwrYQIzxvF8NpY4sg8p2Q,26563
 langtrace_python_sdk/instrumentation/groq/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langtrace_python_sdk/instrumentation/groq/instrumentation.py,sha256=LHzJkICjflnvAhRzEM7s-Pd9tRKlI3qLOYg0_vYDos0,1826
 langtrace_python_sdk/instrumentation/groq/patch.py,sha256=uuEAPZQXGmvuMAb6lIV-1elz9kNe7L3wDwIASUjrQAQ,25572
@@ -92,19 +92,27 @@ langtrace_python_sdk/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJ
 langtrace_python_sdk/utils/llm.py,sha256=KA5ZL397_ktSpAnIOpxYTaSjaaAoxdHYfiAzbbfFfNA,2022
 langtrace_python_sdk/utils/with_root_span.py,sha256=FOiCV59IAhO861ftbwSCIbU9d4a2EBBjxsqc5HPlZf8,2452
 tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-tests/conftest.py,sha256=pgKBuMyWcbzp8zzUPdZRSMj95fJt8AUZf3Gtc1RTWwY,520
-tests/utils.py,sha256=wdbR00LjYfDXzTBFvkksQYEz1hZjGTNKEiw5KPE_bqI,625
-tests/anthropic/conftest.py,sha256=kHa-lSZLGCb-13iUz0JUhFdMnipTwZGswM1alTaNh0Y,776
-tests/anthropic/test_anthropic.py,sha256=tT5H-0JZi9bJtuylFVE0d9B7D15GBvzXTLixcGVCQ0A,4237
+tests/conftest.py,sha256=0Jo6iCZTXbdvyJVhG9UpYGkLabL75378oauCzmt-Sa8,603
+tests/utils.py,sha256=hP8sTH-M8WO6zlLkSFHPf6483ZcKEcnxul6JIIb1pLM,1396
+tests/anthropic/conftest.py,sha256=ht1zF3q4Xb1UiZ3BL_L-CQ3YWrd52Dqb4WTZ3f-GdG4,739
+tests/anthropic/test_anthropic.py,sha256=p4WhcECJ_zctirtkAX6WejrlNfZzo90fSSxqQkit0Ls,3175
 tests/anthropic/cassettes/test_anthropic.yaml,sha256=z7YAqA_BBgI-hw7uyVMdLoIZdBYhKwl9cuTzUu9nAZs,2332
 tests/anthropic/cassettes/test_anthropic_streaming.yaml,sha256=D0w4-6dfsgrhbNLJEj8gZBV0yXfrAfA9u90Yu8Ac-TY,11675
 tests/anthropic/cassettes/test_async_anthropic_streaming.yaml,sha256=hQZPY2vwBaW3BWllLd0lcGQ73DjA8C3Ips1Hx9pA-ao,8373
 tests/chroma/test_chroma.py,sha256=5zeInDcP5VplrM9ex2iuVKcpVKMDfEf_ZDK9D6Tc700,2392
+tests/cohere/conftest.py,sha256=jBbyg-tut-ZJN5_5D64sGmaPIhT_nQQQAiW43kl5Rz4,621
+tests/cohere/test_cohere_chat.py,sha256=MioEIXB5e-4t99Vt-KpLlxd4diPbp3S6eNfnPWjEDXM,4564
+tests/cohere/test_cohere_embed.py,sha256=bFKDzUnypMyLefAOHR-6dEgvt0yys2Yw7-F0zc0uQyc,1250
+tests/cohere/test_cohere_rerank.py,sha256=ZDzvXsY9l1fiytOAKGQidpBkop3fxV6d8xEy_gQt9jE,2279
+tests/cohere/cassettes/test_cohere_chat.yaml,sha256=_Az6Kfhp1XmoK6zwTxqbnRYTmsm0k88McOidwQaRp3g,3436
+tests/cohere/cassettes/test_cohere_chat_streaming.yaml,sha256=TdWe64Pyfa5Ho9lqTMIW41YntUmUwmYQc55qWntbvas,10610
+tests/cohere/cassettes/test_cohere_embed.yaml,sha256=YAxJ3otY9N6jJQOKi50nQA4xY9-dknm1oXEP8KiCH5U,27312
+tests/cohere/cassettes/test_cohere_rerank.yaml,sha256=m5EAKAM5JEcGMyUlmjaxU-cNLiXKcenSFQyU-UaUI0U,2392
 tests/langchain/test_langchain.py,sha256=GGGRcxz0isNmegeum37XFrlJqI6jB6_iUvv8AJ5iG24,2481
 tests/langchain/test_langchain_community.py,sha256=m9lBmMZIeUouKSq1JfdBupV0-0ef39GD6BKsA0Cf_08,2515
 tests/langchain/test_langchain_core.py,sha256=hCuKkIUvDQOUBM2oEgMG3Iq_KNTwC2sH7_Y_IR5FIno,4238
-tests/openai/conftest.py,sha256=G7OeC6qgr1c76S16cuaQrR_eCg6c7qBFKhwKhLnkk10,737
-tests/openai/test_chat_completion.py,sha256=7Rts_mFRWL6NDqQUR1eP1hcgkh63Cpv2neLDqezoyzg,6252
+tests/openai/conftest.py,sha256=BkehS6heg-O91Nzoc4546OSiAzy8KgSgk7VCO3A11zM,700
+tests/openai/test_chat_completion.py,sha256=-TwCAF2czsiGZqETQvMOgmIJgdFWycJ9nOeO5q65AJs,5804
 tests/openai/test_embeddings.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 tests/openai/test_image_generation.py,sha256=0zANtoltorKPel37fSbut-icpJ-wKSN4unoKZOBBsjQ,4397
 tests/openai/cassettes/test_async_chat_completion_streaming.yaml,sha256=0aZHFy9NvXegEDjGWyoG-_ItLr7JYAKbaBLIPSl-pfM,6844
@@ -113,7 +121,7 @@ tests/openai/cassettes/test_chat_completion.yaml,sha256=YkNFgK9VHAzNqGWuxFcTiE19
 tests/openai/cassettes/test_chat_completion_streaming.yaml,sha256=nkx_TemQMYSZxUF_b-LCEFwCRDm0AkQHLf4sdJVuZBw,2592394
 tests/openai/cassettes/test_image_generation.yaml,sha256=gn5aSVp6V6_hb_rt2NnkAWd_idzDxo-7VzhZII0Wslw,3562
 tests/pinecone/test_pinecone.py,sha256=_wlJbSKnY7gyzVcwxIWKft1P_t8dWwcIKNfGCrRLiHs,2633
-langtrace_python_sdk-2.0.1.dist-info/METADATA,sha256=WdTtq8DwBulr_4zE5wAOLCU038RQoLVeglaxmBS4ABk,9585
-langtrace_python_sdk-2.0.1.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
-langtrace_python_sdk-2.0.1.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
-langtrace_python_sdk-2.0.1.dist-info/RECORD,,
+langtrace_python_sdk-2.0.2.dist-info/METADATA,sha256=na35ZlWWnYIFmlXHc02p5Ifk6-2CJ6ZrGDbIM7YYhQ8,9585
+langtrace_python_sdk-2.0.2.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
+langtrace_python_sdk-2.0.2.dist-info/licenses/LICENSE,sha256=QwcOLU5TJoTeUhuIXzhdCEEDDvorGiC6-3YTOl4TecE,11356
+langtrace_python_sdk-2.0.2.dist-info/RECORD,,

tests/anthropic/conftest.py

@@ -33,6 +33,3 @@ def vcr_config():
 @pytest.fixture(scope="session", autouse=True)
 def instrument():
     AnthropicInstrumentation().instrument()
-
-    yield
-    # exporter.shutdown()

tests/anthropic/test_anthropic.py

@@ -2,6 +2,7 @@ import pytest
 import json
 import importlib
 from langtrace_python_sdk.constants.instrumentation.anthropic import APIS
+from tests.utils import assert_response_format, assert_token_count
 
 
 @pytest.mark.vcr()
@@ -16,7 +17,7 @@ def test_anthropic(anthropic_client, exporter):
         "stream": False,
         "max_tokens": 1024,
     }
-    response = anthropic_client.messages.create(**kwargs)
+    anthropic_client.messages.create(**kwargs)
     spans = exporter.get_finished_spans()
     completion_span = spans[-1]
 
@@ -36,20 +37,8 @@ def test_anthropic(anthropic_client, exporter):
     assert attributes.get("llm.prompts") == json.dumps(messages_value)
     assert attributes.get("llm.stream") is False
 
-    tokens = json.loads(attributes.get("llm.token.counts"))
-    output_tokens = tokens.get("output_tokens")
-    prompt_tokens = tokens.get("input_tokens")
-    total_tokens = tokens.get("total_tokens")
-
-    assert output_tokens and prompt_tokens and total_tokens
-    assert output_tokens + prompt_tokens == total_tokens
-
-    langtrace_responses = json.loads(attributes.get("llm.responses"))
-    assert isinstance(langtrace_responses, list)
-    for langtrace_response in langtrace_responses:
-        assert isinstance(langtrace_response, dict)
-        assert "role" in langtrace_response
-        assert "content" in langtrace_response
+    assert_token_count(attributes)
+    assert_response_format(attributes)
 
 
 @pytest.mark.vcr()
@@ -93,17 +82,5 @@ def test_anthropic_streaming(anthropic_client, exporter):
 
     assert len(events) - 2 == chunk_count  # -2 for start and end events
 
-    tokens = json.loads(attributes.get("llm.token.counts"))
-    output_tokens = tokens.get("output_tokens")
-    prompt_tokens = tokens.get("input_tokens")
-    total_tokens = tokens.get("total_tokens")
-
-    assert output_tokens and prompt_tokens and total_tokens
-    assert output_tokens + prompt_tokens == total_tokens
-
-    langtrace_responses = json.loads(attributes.get("llm.responses"))
-    assert isinstance(langtrace_responses, list)
-    for langtrace_response in langtrace_responses:
-        assert isinstance(langtrace_response, dict)
-        assert "role" in langtrace_response
-        assert "content" in langtrace_response
+    assert_token_count(attributes)
+    assert_response_format(attributes)

tests/cohere/cassettes/test_cohere_chat.yaml

@@ -0,0 +1,91 @@
+interactions:
+- request:
+    body: '{"message": "Tell me a story in 3 sentences or less?", "stream": false,
+      "model": "command", "preamble": "answer like a pirate", "chat_history": [{"role":
+      "USER", "message": "Who discovered gravity?"}, {"role": "CHATBOT", "message":
+      "The man who is widely credited with discovering gravity is Sir Isaac Newton"}],
+      "connectors": [{"id": "web-search"}], "temperature": 0.1}'
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '370'
+      content-type:
+      - application/json
+      host:
+      - api.cohere.ai
+      user-agent:
+      - python-httpx/0.27.0
+      x-fern-language:
+      - Python
+      x-fern-sdk-name:
+      - cohere
+      x-fern-sdk-version:
+      - 5.3.2
+    method: POST
+    uri: https://api.cohere.ai/v1/chat
+  response:
+    body:
+      string: '{"response_id":"27cf0e2e-b4d5-427b-9866-d108307022de","text":"Argh,
+        many years ago, a bold explorer named Isaac Newton sat beneath an apple tree
+        for many hours, contemplating the nature of the universe and the motions of
+        the heavenly bodies. He posited that all objects in the universe are attracted
+        to each other. Then one day, his theory was proven true when a juicy apple
+        fell from the tree and struck him on the head, inspiring his theory of gravity
+        and changing science forever.","generation_id":"2a7babaa-f1da-468f-9891-5171f137ba04","chat_history":[{"role":"USER","message":"Who
+        discovered gravity?"},{"role":"CHATBOT","message":"The man who is widely credited
+        with discovering gravity is Sir Isaac Newton"},{"role":"USER","message":"Tell
+        me a story in 3 sentences or less?"},{"role":"CHATBOT","message":"Argh, many
+        years ago, a bold explorer named Isaac Newton sat beneath an apple tree for
+        many hours, contemplating the nature of the universe and the motions of the
+        heavenly bodies. He posited that all objects in the universe are attracted
+        to each other. Then one day, his theory was proven true when a juicy apple
+        fell from the tree and struck him on the head, inspiring his theory of gravity
+        and changing science forever."}],"finish_reason":"COMPLETE","meta":{"api_version":{"version":"1"},"billed_units":{"input_tokens":31,"output_tokens":85},"tokens":{"input_tokens":49,"output_tokens":85}},"documents":[]}'
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      Via:
+      - 1.1 google
+      access-control-expose-headers:
+      - X-Debug-Trace-ID
+      cache-control:
+      - no-cache, no-store, no-transform, must-revalidate, private, max-age=0
+      content-length:
+      - '1420'
+      content-type:
+      - application/json
+      date:
+      - Fri, 26 Apr 2024 13:51:11 GMT
+      expires:
+      - Thu, 01 Jan 1970 00:00:00 UTC
+      num_chars:
+      - '220'
+      num_tokens:
+      - '116'
+      pragma:
+      - no-cache
+      server:
+      - envoy
+      vary:
+      - Origin
+      x-accel-expires:
+      - '0'
+      x-debug-trace-id:
+      - 6e0f777a5f86c8b5453bf374fd3fadbd
+      x-endpoint-monthly-call-limit:
+      - '1000'
+      x-envoy-upstream-service-time:
+      - '3037'
+      x-trial-endpoint-call-limit:
+      - '40'
+      x-trial-endpoint-call-remaining:
+      - '39'
+    status:
+      code: 200
+      message: OK
+version: 1

tests/cohere/cassettes/test_cohere_chat_streaming.yaml

@@ -0,0 +1,275 @@
+interactions:
+- request:
+    body: '{"message": "Tell me a story in 3 sentences or less?", "stream": true,
+      "model": "command", "preamble": "answer like a pirate", "chat_history": [{"role":
+      "USER", "message": "Who discovered gravity?"}, {"role": "CHATBOT", "message":
+      "The man who is widely credited with discovering gravity is Sir Isaac Newton"}],
+      "connectors": [{"id": "web-search"}], "temperature": 0.1}'
+    headers:
+      accept:
+      - '*/*, text/event-stream, application/stream+json'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '369'
+      content-type:
+      - application/json
+      host:
+      - api.cohere.ai
+      user-agent:
+      - python-httpx/0.27.0
+      x-fern-language:
+      - Python
+      x-fern-sdk-name:
+      - cohere
+      x-fern-sdk-version:
+      - 5.3.2
+    method: POST
+    uri: https://api.cohere.ai/v1/chat
+  response:
+    body:
+      string: '{"is_finished":false,"event_type":"stream-start","generation_id":"1d521296-3008-42bf-8062-8b744d997ba0"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"Ar"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"gh"}
+
+        {"is_finished":false,"event_type":"text-generation","text":","}
+
+        {"is_finished":false,"event_type":"text-generation","text":" here"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" be"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" the"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" story"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" of"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" Gold"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"il"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"ocks"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" and"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" the"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" Three"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" Bears"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" told"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" in"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" three"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" sentences"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" or"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" less"}
+
+        {"is_finished":false,"event_type":"text-generation","text":":"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"\n\nOnce"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" upon"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" a"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" time"}
+
+        {"is_finished":false,"event_type":"text-generation","text":","}
+
+        {"is_finished":false,"event_type":"text-generation","text":" a"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" little"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" girl"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" named"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" Gold"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"il"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"ocks"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" stumbled"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" into"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" the"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" home"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" of"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" three"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" bears"}
+
+        {"is_finished":false,"event_type":"text-generation","text":","}
+
+        {"is_finished":false,"event_type":"text-generation","text":" she"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" ate"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" their"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" food"}
+
+        {"is_finished":false,"event_type":"text-generation","text":","}
+
+        {"is_finished":false,"event_type":"text-generation","text":" broke"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" their"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" chairs"}
+
+        {"is_finished":false,"event_type":"text-generation","text":","}
+
+        {"is_finished":false,"event_type":"text-generation","text":" and"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" slept"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" in"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" their"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" beds"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"."}
+
+        {"is_finished":false,"event_type":"text-generation","text":" When"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" the"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" three"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" bears"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" returned"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" they"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" were"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" outraged"}
+
+        {"is_finished":false,"event_type":"text-generation","text":","}
+
+        {"is_finished":false,"event_type":"text-generation","text":" but"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" Gold"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"il"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"ocks"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" had"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" disappeared"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" -"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" she"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" had"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" run"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" away"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"."}
+
+        {"is_finished":false,"event_type":"text-generation","text":"\n\nAr"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"gh"}
+
+        {"is_finished":false,"event_type":"text-generation","text":","}
+
+        {"is_finished":false,"event_type":"text-generation","text":" that"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" be"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" the"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" end"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" of"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" the"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" story"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" me"}
+
+        {"is_finished":false,"event_type":"text-generation","text":" heart"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"ies"}
+
+        {"is_finished":false,"event_type":"text-generation","text":"!"}
+
+        {"is_finished":true,"event_type":"stream-end","response":{"response_id":"86b73a20-0cfa-414b-b148-e7857669ef3e","text":"Argh,
+        here be the story of Goldilocks and the Three Bears told in three sentences
+        or less:\n\nOnce upon a time, a little girl named Goldilocks stumbled into
+        the home of three bears, she ate their food, broke their chairs, and slept
+        in their beds. When the three bears returned they were outraged, but Goldilocks
+        had disappeared - she had run away.\n\nArgh, that be the end of the story
+        me hearties!","generation_id":"1d521296-3008-42bf-8062-8b744d997ba0","chat_history":[{"role":"USER","message":"Who
+        discovered gravity?"},{"role":"CHATBOT","message":"The man who is widely credited
+        with discovering gravity is Sir Isaac Newton"},{"role":"USER","message":"Tell
+        me a story in 3 sentences or less?"},{"role":"CHATBOT","message":"Argh, here
+        be the story of Goldilocks and the Three Bears told in three sentences or
+        less:\n\nOnce upon a time, a little girl named Goldilocks stumbled into the
+        home of three bears, she ate their food, broke their chairs, and slept in
+        their beds. When the three bears returned they were outraged, but Goldilocks
+        had disappeared - she had run away.\n\nArgh, that be the end of the story
+        me hearties!"}],"finish_reason":"COMPLETE","meta":{"api_version":{"version":"1"},"billed_units":{"input_tokens":31,"output_tokens":96},"tokens":{"input_tokens":49,"output_tokens":96}},"documents":[]},"finish_reason":"COMPLETE"}
+
+        '
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      Via:
+      - 1.1 google
+      access-control-expose-headers:
+      - X-Debug-Trace-ID
+      cache-control:
+      - no-cache, no-store, no-transform, must-revalidate, private, max-age=0
+      content-type:
+      - application/stream+json
+      date:
+      - Fri, 26 Apr 2024 13:59:03 GMT
+      expires:
+      - Thu, 01 Jan 1970 00:00:00 UTC
+      pragma:
+      - no-cache
+      server:
+      - envoy
+      transfer-encoding:
+      - chunked
+      vary:
+      - Origin
+      x-accel-expires:
+      - '0'
+      x-debug-trace-id:
+      - 2803337aa161b7c4f034921bb2bb1444
+      x-endpoint-monthly-call-limit:
+      - '1000'
+      x-envoy-upstream-service-time:
+      - '66'
+      x-trial-endpoint-call-limit:
+      - '40'
+      x-trial-endpoint-call-remaining:
+      - '39'
+    status:
+      code: 200
+      message: OK
+version: 1

tests/cohere/cassettes/test_cohere_embed.yaml

@@ -0,0 +1,73 @@
+interactions:
+- request:
+    body: '{"texts": ["hello", "goodbye"], "model": "embed-english-v3.0", "input_type":
+      "classification"}'
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '94'
+      content-type:
+      - application/json
+      host:
+      - api.cohere.ai
+      user-agent:
+      - python-httpx/0.27.0
+      x-fern-language:
+      - Python
+      x-fern-sdk-name:
+      - cohere
+      x-fern-sdk-version:
+      - 5.3.2
+    method: POST
+    uri: https://api.cohere.ai/v1/embed
+  response:
+    body:
+      string: '{"id":"2c1ede34-c4c7-4cdd-a74e-ce035af81e20","texts":["hello","goodbye"],"embeddings":[[0.016296387,-0.008354187,-0.04699707,-0.07104492,0.00013196468,-0.014892578,-0.018661499,0.019134521,0.008476257,0.04159546,-0.036895752,-0.00048303604,0.06414795,-0.036346436,0.045806885,-0.03125,0.03793335,0.048583984,0.0062179565,0.0071144104,-0.020935059,0.04196167,-0.039398193,0.03463745,0.051879883,0.030838013,-0.0048103333,-0.00036287308,-0.017944336,-0.039611816,0.013389587,0.0044021606,0.018951416,0.020767212,-0.0025997162,0.0904541,-0.0121154785,-0.026184082,0.012413025,0.004119873,0.030654907,-0.030792236,-0.041107178,-0.02368164,-0.043304443,-0.00077438354,-0.017074585,-0.019729614,0.078125,-0.031585693,0.020217896,-0.01524353,0.017471313,-0.0008010864,-0.03717041,0.011062622,-0.072143555,-0.013175964,0.01058197,0.030853271,0.044799805,0.0045928955,0.03253174,0.047698975,-0.0039024353,-0.01965332,0.024475098,-0.013755798,0.018951416,-0.015487671,0.015594482,0.00096321106,-0.006450653,-0.04748535,-0.021972656,0.06323242,-0.009498596,0.014297485,0.0038471222,-0.023117065,-0.02180481,-0.01928711,-0.08758545,-0.04852295,0.029510498,0.011276245,-0.013504028,-0.009391785,-0.0064468384,0.010978699,-0.014404297,0.053741455,0.046569824,0.00042700768,-0.037719727,0.011985779,-0.009643555,0.0067749023,0.008071899,0.018829346,-0.05419922,-0.020950317,-0.02659607,-0.028869629,-0.015716553,0.022705078,-0.0046958923,0.02192688,0.032440186,0.048034668,-0.006843567,0.045074463,-0.02293396,0.010238647,-0.04534912,0.01638794,-0.00680542,0.0038871765,-0.032836914,0.051361084,0.0395813,0.032928467,-0.00843811,0.007858276,-0.040802002,-0.008346558,-0.013252258,-0.046173096,0.051727295,-0.027175903,-0.011497498,0.04940796,-0.095214844,-0.0345459,-0.021453857,0.0051002502,-0.01725769,-0.045196533,-0.0016956329,0.021575928,0.07720947,-0.00094270706,0.020904541,0.05001831,-0.033111572,0.032287598,-0.0052833557,-0.00007402897,0.035125732,0.019424438,-0.06665039,-0.02557373,0.010887146,0.05807495,0.015022278,0.0657959,-0.015350342,0.008468628,-0.017944336,0.029388428,-0.005126953,0.015914917,0.051879883,-0.015975952,-0.039031982,-0.012374878,0.0032424927,0.0008568764,0.014579773,0.021530151,-0.0061912537,0.028717041,0.046844482,0.032836914,0.0071372986,-0.023406982,-0.03717041,0.016723633,0.03994751,0.025390625,0.03427124,-0.01914978,-0.026000977,0.07342529,-0.03213501,-0.058258057,0.029144287,0.001042366,0.030517578,0.011474609,0.058410645,0.005027771,-0.038635254,-0.015029907,-0.015655518,-0.03918457,-0.016342163,-0.020858765,-0.0043907166,0.03857422,0.007423401,-0.0473938,0.04257202,-0.043823242,-0.03842163,-0.033691406,-0.010925293,0.012260437,0.0009822845,0.0058937073,-0.008644104,-0.031585693,0.0055618286,-0.06976318,-0.030578613,-0.038970947,-0.08880615,-0.00315094,0.00020766258,0.04058838,0.0028266907,-0.0018129349,-0.01625061,-0.022277832,-0.008956909,-0.009292603,-0.040771484,-0.008705139,-0.065979004,-0.010414124,-0.0152282715,0.033447266,-0.033599854,-0.008049011,-0.020828247,0.0053901672,0.0002875328,0.037078857,0.015159607,-0.0016326904,0.012397766,0.0026817322,-0.032196045,-0.0079422,0.03567505,-0.0010242462,0.03652954,-0.0035171509,0.01802063,0.026641846,0.0107421875,-0.021942139,0.035095215,-0.0236969,-0.015975952,0.039215088,0.0038166046,0.020462036,-0.039764404,0.035888672,-0.038604736,-0.008621216,-0.012619019,-0.014602661,-0.036102295,-0.02368164,-0.0121536255,-0.0054512024,-0.015701294,-0.016296387,0.016433716,-0.005672455,-0.019332886,0.00025129318,0.0803833,0.04248047,-0.05960083,-0.009147644,-0.0021247864,0.012481689,-0.015129089,-0.021133423,-0.01878357,0.0027332306,0.036956787,-0.0053253174,-0.0007238388,0.016983032,-0.0034694672,0.059387207,0.076660156,0.015312195,-0.015823364,0.02456665,0.012901306,0.020126343,-0.032440186,0.011291504,-0.001876831,-0.052215576,0.004634857,0.036956787,0.006164551,-0.023422241,-0.025619507,0.024261475,0.023849487,0.015007019,0.020050049,-0.044067383,0.030029297,0.021377563,0.011657715,0.017196655,-0.032318115,-0.031555176,-0.00982666,-0.0039787292,-0.079589844,-0.006416321,0.00844574,-0.007434845,-0.045013428,-0.02557373,-0.01537323,0.027633667,-0.076538086,-0.0025749207,-0.05279541,0.029373169,0.047912598,0.00083875656,-0.01234436,-0.017059326,0.01159668,0.014228821,0.029571533,-0.055114746,0.006389618,0.028869629,0.09375,-0.014251709,0.029418945,0.007633209,0.010848999,-0.004055023,-0.02116394,0.007194519,-0.0062217712,-0.01209259,0.024749756,-0.037506104,-0.029510498,-0.028442383,0.03189087,0.0008239746,0.007419586,-0.016723633,0.06964111,-0.07232666,0.022201538,-0.019882202,-0.0385437,-0.022567749,0.010353088,-0.027755737,-0.006713867,-0.023406982,-0.025054932,-0.013076782,0.015808105,-0.0073165894,0.02949524,-0.036499023,-0.07287598,-0.01876831,-0.02709961,-0.06567383,0.050567627,0.004047394,0.030471802,0.025405884,0.046783447,0.01763916,0.053466797,0.049072266,-0.015197754,0.0013389587,0.049591064,0.006965637,-0.00014233589,0.01335907,-0.04675293,-0.026733398,0.03024292,0.0012464523,-0.037200928,0.030166626,-0.08544922,-0.013893127,-0.014823914,0.0014219284,-0.023620605,-0.0010480881,-0.072387695,0.057922363,-0.04067993,-0.025299072,0.020446777,0.06451416,0.007205963,0.015838623,-0.008674622,0.0002270937,-0.026321411,0.027130127,-0.01828003,-0.011482239,0.03463745,0.00724411,-0.010406494,0.025268555,-0.023651123,0.04034424,-0.036834717,0.05014038,-0.026184082,0.036376953,0.03253174,-0.01828003,-0.023376465,-0.034576416,-0.00598526,-0.023239136,-0.032409668,0.07672119,-0.038604736,0.056884766,-0.012550354,-0.03778076,-0.013061523,0.017105103,0.010482788,-0.005077362,-0.010719299,-0.018661499,0.019760132,0.022018433,-0.058746338,0.03564453,-0.0892334,0.025421143,-0.015716553,0.07910156,-0.009361267,0.016921997,0.048736572,0.035247803,0.01864624,0.011413574,0.018295288,0.00052690506,-0.07122803,-0.01890564,-0.017669678,0.027694702,0.0152282715,0.006511688,-0.045837402,-0.009765625,0.013877869,-0.0146102905,0.033294678,-0.0019874573,0.023040771,0.025619507,-0.015823364,-0.020858765,-0.023529053,0.0070152283,-0.0647583,0.036224365,0.0023403168,-0.062286377,-0.036315918,0.021209717,-0.037353516,-0.03656006,0.01889038,0.023239136,0.011764526,0.005970001,0.049346924,-0.006893158,-0.015068054,-0.0008716583,-0.0034999847,0.04034424,0.017913818,-0.06707764,-0.07531738,0.00042319298,-0.00680542,-0.0023174286,0.04425049,-0.05105591,-0.016967773,0.020507812,0.038604736,0.029846191,0.04309082,-0.00084733963,-0.008911133,0.0082092285,-0.0050239563,0.05038452,0.014595032,0.015182495,0.007247925,-0.04046631,-0.011169434,-0.010292053,0.068603516,0.02470398,-0.0023403168,0.005996704,-0.0010709763,0.008178711,-0.029205322,-0.025253296,0.05822754,0.04269409,0.059295654,-0.0011911392,-0.031311035,0.023712158,-0.037506104,0.004589081,0.014923096,-0.019866943,-0.019180298,-0.0020999908,-0.008972168,0.01348114,0.014801025,-0.02645874,0.019897461,0.081970215,-0.05822754,0.09399414,0.001209259,-0.050750732,0.062316895,-0.014892578,-0.019104004,-0.036987305,-0.040618896,-0.008163452,-0.0035247803,0.06774902,-0.001420021,-0.0013103485,-0.031799316,-0.0023651123,0.012298584,0.003583908,0.050964355,-0.01802063,-0.007091522,0.01448822,-0.016159058,-0.019439697,-0.022491455,-0.036346436,-0.03491211,-0.0032920837,0.003528595,-0.0016469955,0.01612854,-0.003709793,0.012840271,0.0043182373,-0.030456543,0.007369995,0.0039787292,0.036499023,0.021362305,0.00062942505,0.0047073364,0.026382446,-0.0020542145,-0.038757324,-0.00095272064,0.0019435883,0.007232666,-0.0031471252,0.019943237,-0.062042236,0.010826111,0.0026607513,-0.04727173,0.020126343,0.046417236,-0.03881836,0.011222839,0.011428833,-0.056396484,0.010879517,-0.011772156,-0.0038414001,0.010246277,-0.020141602,-0.011169434,0.006916046,-0.022659302,0.010299683,0.046966553,0.0234375,-0.0016288757,-0.03262329,-0.01689148,-0.00031924248,0.028152466,0.004234314,0.03878784,-0.03579712,0.007457733,-0.0036907196,0.0073051453,-0.00028276443,-0.0067100525,0.003206253,-0.0021209717,-0.05960083,0.024337769,0.076171875,-0.012062073,-0.0032787323,-0.08380127,0.024917603,0.019073486,-0.012031555,-0.03237915,-0.0042686462,-0.01525116,-0.0158844,-0.0014514923,-0.024429321,-0.028442383,0.020843506,0.007133484,0.024230957,0.0002002716,-0.005466461,-0.0032367706,0.012718201,0.032806396,0.062042236,-0.040283203,-0.025497437,0.045013428,0.054473877,-0.033599854,-0.0039482117,0.02268982,-0.0012645721,0.045166016,0.0501709,-0.0022602081,0.019897461,0.007926941,0.017364502,0.011650085,-0.042510986,-0.059448242,0.030014038,0.039611816,0.015571594,0.04031372,-0.0006723404,-0.03353882,-0.05569458,0.040283203,0.019058228,-0.032592773,0.004470825,0.06359863,0.029693604,0.01826477,-0.0104522705,-0.043945312,-0.01802063,0.0075187683,-0.02456665,0.02798462,0.0047340393,-0.017623901,-0.014335632,-0.04550171,-0.0039711,0.023864746,-0.015281677,0.055755615,-0.04864502,0.033599854,0.024810791,-0.03048706,-0.043121338,0.011291504,0.024932861,-0.0020275116,0.032287598,-0.0234375,0.006942749,-0.007221222,-0.03869629,-0.03765869,-0.03475952,-0.046936035,0.03012085,-0.021362305,-0.023452759,0.051239014,-0.009925842,0.04925537,-0.00944519,-0.040008545,-0.019485474,-0.00022566319,-0.017028809,0.03277588,0.0066375732,-0.013328552,0.01864624,-0.011726379,0.023849487,0.04006958,0.03793335,0.060821533,0.005504608,-0.0395813,-0.010131836,0.046539307,0.030136108,0.002231598,0.042236328,0.014755249,0.047058105,-0.017318726,0.008598328,0.01966858,0.0064430237,0.03616333,-0.011985779,-0.003446579,-0.06616211,-0.0657959,0.014137268,0.044677734,-0.03515625,-0.05215454,-0.012710571,0.0047416687,0.05368042,0.013900757,0.05001831,0.027709961,0.02557373,-0.025512695,0.0031032562,0.072143555,0.018829346,0.0073928833,0.009269714,-0.011299133,0.0048828125,0.014808655,-0.0184021,-0.00089359283,-0.0015716553,-0.012863159,0.0074386597,-0.020767212,0.02204895,-0.027404785,-0.021972656,0.02494812,0.044006348,-0.011581421,0.06298828,0.009010315,0.03842163,-0.00005555153,0.06774902,0.036254883,-0.016311646,-0.000004887581,0.0057373047,0.03704834,-0.041503906,0.0074043274,-0.012290955,-0.020263672,-0.0057792664,-0.025878906,-0.021652222,-0.008079529,0.022613525,-0.012069702,0.050079346,-0.004283905,-0.021118164,-0.010559082,-0.0041160583,-0.00026345253,-0.01260376,0.050628662,-0.03137207,0.027526855,-0.052642822,-0.0046463013,0.04937744,-0.0017156601,0.014625549,-0.022476196,0.02571106,0.043884277,-0.016952515,-0.021011353,0.056396484,0.056762695,0.013473511,-0.02357483,0.043792725,0.032470703,-0.052612305,-0.017837524,-0.000067055225,0.039276123,-0.012283325,-0.0029888153,-0.024719238,0.012870789,-0.032287598,0.028839111,0.008056641,0.011100769,-0.034210205,0.028198242,0.01940918,0.029052734,0.030303955,0.03475952,-0.03982544,0.026870728,0.02079773,0.03012085,-0.044281006,0.006462097,-0.008705139,-0.024734497,0.02458191,-0.050201416,-0.028778076,0.036956787,0.025634766,-0.025650024,0.020629883,-0.04385376,0.009536743,-0.0027256012,0.031158447,0.008712769,-0.039855957,-0.018249512,-0.011268616,0.009689331,-0.032073975,0.023010254,0.04925537,0.013168335,0.02734375,0.031707764,-0.024032593,-0.010604858,-0.00258255,0.0054092407,0.033569336,0.0068359375,0.019882202,0.018096924,-0.05392456,-0.0030059814,-0.01374054,-0.008483887,0.016494751,-0.015487671,0.016143799,-0.028198242,-0.016326904,-0.013160706,-0.046905518,0.026428223,-0.02420044,-0.022262573,0.041748047,0.05557251,-0.0044059753,-0.030960083,-0.023544312,0.0103302,-0.013534546,-0.016830444,0.028167725,0.0061950684,0.02178955,-0.06945801,-0.040039062,-0.0024642944,-0.06359863,-0.020812988,0.029006958,0.0072364807,-0.028747559,-0.057891846,0.022155762,-0.035369873,-0.025909424,-0.04095459,0.0019893646,-0.0038146973,-0.030639648,-0.038970947,-0.0026626587,-0.0047454834,-0.014816284,0.008575439,-0.032165527,-0.011062622,0.003622055,-0.0129852295,-0.0007658005,-0.009902954,0.03704834,-0.02456665,0.020385742,0.0019044876,-0.008552551,-0.028137207,-0.006500244,0.017227173,-0.0077285767,-0.05496216,0.038024902,-0.0335083,0.047668457,-0.02998352,-0.0395813,-0.0068359375,-0.024627686,-0.005756378,0.025863647,0.032104492,-0.029022217,-0.08685303,-0.014724731,-0.035583496,0.024002075,0.008422852,0.012931824,-0.0055656433,-0.013748169,-0.021530151,-0.034332275,-0.008766174,-0.025222778,0.019836426,-0.011619568,-0.037963867,0.013519287,-0.035736084,0.049102783,-0.011398315,0.050598145,-0.066833496,0.080566406,-0.061553955,-0.041778564,0.01864624,0.014907837,-0.010482788,0.035217285,-0.0473938,-0.031951904,0.052886963,-0.022109985,0.031677246,-0.01977539,0.08282471,0.012901306,-0.009490967,0.0030956268,0.023895264,0.012611389,-0.0011844635,-0.007633209,0.019195557,-0.05404663,0.006187439,-0.06762695,-0.049468994,0.028121948,-0.004032135,-0.043151855,0.028121948,-0.0058555603,0.019454956,0.0028438568,-0.0036354065,-0.015411377,-0.026535034,0.03704834,-0.01802063,0.009765625],[0.04663086,-0.023239136,0.008163452,-0.03945923,-0.018051147,-0.011123657,0.0022335052,-0.0015516281,-0.002336502,0.031799316,-0.049591064,-0.049835205,0.019317627,-0.013328552,-0.01838684,-0.067871094,0.02671814,0.038085938,0.03265381,-0.0043907166,0.026321411,0.0070114136,-0.037628174,0.008026123,0.015525818,0.066589355,-0.018005371,-0.0017309189,-0.052368164,-0.055511475,-0.00504303,0.043029785,-0.013328552,0.08581543,-0.038269043,0.051971436,-0.04675293,0.038146973,0.05328369,-0.028762817,0.01625061,-0.008644104,-0.060150146,-0.0259552,-0.05432129,-0.00680542,-0.012649536,0.0025501251,0.060272217,-0.013168335,0.046691895,0.030395508,0.039733887,0.00044679642,-0.034240723,0.01828003,-0.047546387,-0.036499023,0.024505615,0.027374268,0.015197754,-0.003932953,0.03475952,0.013633728,0.020858765,-0.025344849,-0.056732178,0.008178711,0.043304443,0.014625549,-0.0020503998,-0.033569336,-0.00178051,-0.0446167,-0.045837402,0.089538574,0.00440979,0.03741455,0.0015287399,-0.035339355,0.017654419,-0.008956909,-0.035064697,-0.014251709,0.008331299,0.0077781677,0.0020999908,-0.021636963,-0.014625549,-0.0209198,-0.009429932,0.070617676,0.013923645,-0.025558472,-0.0519104,-0.0049552917,0.000998497,-0.01448822,-0.027175903,-0.04083252,-0.032043457,-0.0096588135,-0.047088623,-0.0012331009,-0.025878906,0.031799316,-0.023712158,0.015701294,0.017730713,0.062927246,0.009178162,-0.046295166,-0.014701843,-0.007751465,-0.021148682,0.033966064,-0.013664246,0.03945923,-0.02520752,0.08905029,-0.039520264,-0.012435913,-0.057403564,0.007068634,0.006061554,-0.040161133,-0.015548706,0.080078125,0.08862305,0.008003235,-0.048339844,0.037750244,-0.04498291,-0.065979004,-0.032470703,-0.03225708,0.004890442,-0.013023376,-0.020965576,0.035095215,0.035491943,-0.01486969,0.027023315,0.009552002,-0.01285553,0.044891357,0.00062322617,-0.030639648,0.024108887,0.0035648346,-0.06585693,-0.011070251,0.037506104,0.05697632,-0.027236938,0.03475952,0.0143585205,-0.014442444,-0.011405945,-0.013648987,-0.028625488,0.024902344,0.09387207,-0.012741089,-0.040985107,-0.018814087,0.0046920776,-0.017715454,0.013839722,0.0022621155,0.0024433136,-0.028366089,-0.0046310425,0.028717041,-0.00013160706,0.006690979,-0.053863525,0.03302002,0.040802002,0.03201294,0.032073975,-0.03125,-0.005241394,0.048828125,-0.016204834,-0.0014667511,-0.013572693,0.007949829,0.019744873,-0.004776001,-0.0022506714,0.033111572,0.00039958954,0.008369446,-0.021057129,-0.033935547,-0.03692627,0.0042762756,-0.030380249,-0.01876831,-0.023529053,0.004764557,0.026947021,-0.013267517,-0.023666382,0.0024929047,-0.017990112,0.035217285,0.0034389496,0.030380249,0.02015686,-0.013061523,-0.047790527,0.042633057,0.009559631,-0.03186035,-0.02796936,-0.0151901245,-0.0039482117,0.0345459,-0.018096924,0.012062073,-0.02180481,0.031402588,0.041412354,-0.052459717,0.006286621,-0.033203125,-0.0013237,-0.012466431,-0.041748047,0.027313232,-0.0284729,-0.05682373,-0.02809143,0.030899048,0.023773193,0.044677734,-0.0064353943,-0.0000064373016,0.011512756,0.0028190613,-0.041870117,-0.028182983,0.014595032,-0.0143966675,0.022949219,-0.004371643,0.01461792,0.0035171509,0.01398468,-0.04473877,0.04232788,-0.033599854,-0.000647068,0.034606934,0.006160736,-0.014640808,0.028137207,-0.02470398,0.0043563843,0.00039553642,-0.039886475,0.014251709,-0.035736084,-0.021347046,-0.029663086,-0.011688232,-0.038085938,-0.0034008026,0.029144287,-0.010948181,-0.024978638,0.009468079,0.093933105,0.014205933,-0.08569336,-0.011657715,0.02027893,0.0063095093,-0.0035533905,0.020446777,0.029968262,-0.002008438,0.03253174,0.029891968,0.019577026,-0.002922058,-0.009994507,0.029418945,0.049987793,0.046295166,-0.0072898865,0.019638062,0.042816162,0.0066108704,0.06591797,0.04714966,-0.026062012,-0.019470215,0.009979248,0.018081665,0.000009059906,-0.043060303,-0.0043907166,0.064331055,0.051605225,-0.0040893555,0.018081665,-0.024749756,-0.014915466,-0.048614502,0.023483276,0.013282776,-0.011741638,-0.036346436,-0.0076293945,0.023086548,-0.051849365,0.023223877,0.033721924,-0.003929138,-0.044647217,0.020019531,-0.029678345,-0.0031986237,0.030548096,-0.040161133,-0.020874023,0.028793335,0.037872314,0.011314392,-0.030838013,-0.051818848,-0.007774353,0.0070724487,0.02507019,-0.0112838745,0.014930725,0.010543823,0.085998535,0.019332886,0.0107803345,0.00014901161,0.001613617,-0.024993896,-0.04940796,0.010643005,0.04269409,-0.02571106,0.001124382,-0.018844604,-0.014953613,0.027786255,0.033447266,0.0038719177,0.011268616,0.004295349,0.028656006,-0.078063965,-0.012619019,-0.03527832,-0.061279297,0.0625,0.038116455,-0.008308411,-0.017913818,0.031311035,-0.018722534,0.0362854,-0.019363403,0.021362305,-0.0029010773,-0.030288696,-0.07293701,0.008544922,0.006755829,-0.068237305,0.0491333,0.016494751,-0.021621704,0.020980835,0.026443481,0.051879883,0.035583496,0.030548096,-0.03366089,-0.017532349,0.066101074,0.03930664,0.013633728,-0.008621216,0.031982422,-0.042388916,-0.00042247772,-0.020492554,0.04006958,0.052825928,-0.0044136047,-0.02243042,-0.04260254,0.02418518,-0.020584106,-0.0027770996,-0.05908203,0.026611328,-0.046051025,-0.03451538,0.017944336,0.054260254,0.019348145,0.0070114136,0.014205933,-0.019454956,-0.021514893,0.010383606,0.050109863,0.020584106,-0.031677246,-0.048187256,0.01449585,0.04650879,0.025222778,0.004135132,0.02017212,0.044311523,-0.03427124,-0.023757935,0.03479004,-0.012031555,-0.030380249,-0.021560669,-0.010375977,-0.05041504,-0.060821533,0.012283325,-0.026367188,0.061920166,0.026367188,-0.037078857,-0.015136719,0.033355713,-0.010055542,0.025314331,-0.027893066,-0.010032654,0.017684937,-0.00002783537,-0.061157227,0.030273438,-0.103759766,0.035583496,-0.028167725,0.07171631,-0.0211792,-0.013725281,0.04437256,0.041137695,0.027145386,0.032073975,0.008926392,-0.021560669,0.007381439,0.019165039,0.0012969971,-0.01928711,0.026672363,-0.01222229,-0.056365967,0.010398865,-0.02255249,0.00093221664,-0.009353638,0.016082764,0.022872925,0.025024414,-0.024459839,0.040618896,-0.049224854,-0.0035133362,-0.047698975,0.01727295,0.034057617,-0.004096985,-0.009361267,0.011291504,-0.010093689,-0.017990112,0.04107666,-0.058563232,-0.03387451,-0.046905518,0.015411377,-0.02003479,-0.010528564,-0.01689148,0.010391235,-0.040618896,0.029205322,-0.020492554,-0.082092285,0.0004811287,0.043518066,-0.044830322,0.020141602,-0.02319336,0.0024662018,0.012825012,0.04977417,0.06225586,0.027801514,0.005153656,0.04147339,0.0011873245,0.004486084,-0.02494812,0.061706543,0.012184143,-0.0027637482,-0.018447876,-0.008987427,-0.0362854,0.10205078,0.026138306,-0.056549072,0.015899658,0.04449463,-0.017837524,-0.0044898987,-0.04348755,0.06689453,0.008728027,0.047454834,0.03289795,-0.034851074,0.04675293,-0.058807373,0.03164673,0.01322937,-0.06958008,-0.042816162,-0.022918701,-0.019760132,0.008293152,0.02709961,-0.05822754,0.011459351,-0.0008597374,-0.01574707,0.027954102,-0.029785156,-0.03665161,0.017562866,-0.027297974,-0.024017334,-0.0423584,-0.039245605,0.0028457642,-0.0010719299,0.01763916,0.009902954,-0.023849487,-0.009399414,-0.016464233,0.045074463,-0.0056762695,0.04537964,-0.04397583,-0.025817871,0.037353516,-0.018737793,0.01084137,0.0038528442,-0.04547119,-0.024475098,-0.05545044,-0.005756378,0.008132935,0.014541626,-0.0020751953,0.03793335,-0.004421234,-0.037261963,-0.00818634,0.026733398,0.04776001,-0.012313843,0.0019369125,-0.0006084442,0.01335907,-0.033813477,-0.024459839,0.046783447,-0.006389618,-0.055999756,-0.059295654,0.008743286,-0.033966064,0.022537231,-0.018722534,-0.041259766,0.040039062,0.028747559,-0.03515625,0.0019016266,0.041778564,-0.0046539307,0.00014257431,0.011451721,0.016998291,0.00522995,-0.04837036,-0.024520874,0.025466919,-0.020706177,0.017608643,0.062042236,-0.0039596558,-0.021911621,-0.013893127,-0.0000885129,0.00075626373,0.03414917,0.011314392,0.018661499,-0.009719849,0.012748718,-0.026809692,-0.01436615,0.021469116,-0.036254883,0.00907135,-0.026016235,-0.01625061,0.030075073,0.011817932,-0.0038528442,-0.0028858185,-0.021820068,0.037475586,0.0115356445,-0.0077285767,-0.05328369,-0.051361084,0.040649414,-0.005958557,-0.02279663,0.01953125,-0.016937256,0.03781128,-0.0016212463,0.015098572,-0.01626587,0.0067443848,0.027175903,0.011459351,0.038513184,0.06222534,-0.0073547363,-0.010383606,0.0017681122,0.045043945,-0.044921875,-0.0104599,0.035858154,-0.008323669,0.0025901794,0.021514893,-0.010971069,0.016738892,0.0018157959,-0.0071258545,-0.029022217,-0.047027588,-0.02670288,0.029220581,-0.022750854,0.025054932,-0.008544922,0.006164551,-0.029052734,-0.031066895,0.06304932,-0.044647217,-0.017562866,-0.0068511963,0.06604004,0.039916992,-0.007041931,-0.02772522,-0.05795288,-0.022247314,-0.02810669,-0.03845215,0.045074463,-0.014060974,-0.016174316,0.046722412,-0.0006046295,-0.019500732,-0.025985718,0.032989502,0.028366089,0.0021324158,0.0020503998,0.051574707,0.009117126,-0.03112793,-0.006565094,0.019226074,0.009971619,-0.0064735413,-0.017700195,0.0024414062,-0.0008454323,-0.04071045,-0.034820557,-0.031066895,-0.044677734,0.039398193,-0.012580872,-0.06549072,0.027130127,-0.0309906,0.023727417,-0.019760132,0.0066490173,-0.004798889,0.009155273,-0.009902954,0.047576904,0.005466461,0.001537323,0.014862061,-0.0027828217,-0.0079956055,0.043182373,0.0051841736,0.034484863,-0.028015137,-0.012870789,-0.019714355,0.036071777,0.015716553,-0.016860962,0.0034122467,-0.014289856,0.039031982,0.017730713,-0.013549805,0.046691895,0.022094727,0.04647827,0.008033752,0.028747559,-0.030288696,-0.018722534,-0.015113831,0.051971436,-0.040893555,-0.039978027,-0.0042266846,-0.008346558,0.059814453,0.0011167526,0.056030273,-0.08166504,-0.059631348,-0.015731812,0.009529114,0.025756836,0.022232056,-0.0049819946,0.021118164,-0.020446777,0.0032253265,0.017105103,-0.030944824,0.010154724,-0.021881104,-0.018081665,0.029342651,0.024047852,0.017700195,-0.02268982,0.018356323,0.026519775,0.032226562,-0.004711151,0.018753052,0.007789612,0.033172607,-0.034423828,0.035247803,-0.019729614,-0.021194458,0.0071411133,-0.014549255,-0.0073165894,-0.05596924,0.015060425,-0.014305115,-0.030090332,0.001613617,-0.026809692,-0.02571106,-0.0041275024,0.027389526,-0.0059509277,0.0473938,-0.0002002716,0.00037145615,0.0031642914,-0.0044441223,0.0023765564,0.0121154785,0.04260254,-0.035736084,0.019424438,-0.005558014,0.0038166046,0.03717041,-0.0031261444,0.0446167,0.015098572,-0.0022087097,0.0385437,0.024505615,-0.03353882,-0.028533936,0.06048584,-0.019332886,-0.046539307,0.007232666,-0.031585693,0.02168274,0.0046195984,-0.041412354,0.032592773,0.056671143,0.031173706,-0.011398315,0.033416748,0.01802063,-0.0259552,-0.0028705597,0.046539307,-0.040008545,0.022567749,0.020980835,0.024383545,0.02861023,0.010574341,-0.008300781,0.024261475,0.030319214,-0.011238098,-0.030197144,0.013389587,0.010879517,-0.031311035,0.035308838,-0.014755249,0.01612854,0.05722046,-0.019470215,-0.014045715,0.022842407,-0.085998535,0.017166138,0.011474609,0.018325806,0.010398865,0.00434494,-0.013153076,0.025482178,0.007217407,-0.0017223358,0.041046143,0.036895752,-0.028656006,-0.008026123,0.026550293,-0.0146102905,0.0053215027,-0.057037354,0.008743286,0.018066406,0.0025310516,-0.0035171509,-0.02230835,-0.018218994,0.0069618225,-0.006111145,0.017532349,0.034210205,-0.040496826,0.031433105,-0.006587982,-0.031097412,-0.0154418945,-0.009414673,0.006729126,0.004711151,0.00920105,0.0025501251,-0.0016479492,-0.0107803345,-0.070129395,-0.046203613,0.06616211,-0.019622803,-0.06298828,-0.022628784,0.04156494,0.026672363,-0.11505127,-0.080200195,-0.0491333,-0.03744507,-0.0178833,0.016326904,0.03201294,-0.013259888,-0.042114258,0.0023727417,0.005683899,-0.027908325,0.040039062,-0.055847168,-0.03781128,-0.018753052,0.03274536,0.0121536255,0.04360962,-0.0110321045,0.017913818,-0.0231781,-0.018936157,-0.002658844,0.011222839,-0.0082473755,-0.0039043427,0.011512756,-0.014328003,0.037994385,-0.020767212,0.025314331,-0.023727417,0.030303955,0.03302002,0.0040512085,-0.074401855,0.027450562,-0.030838013,0.042053223,-0.04425049,-0.022613525,0.0025463104,0.029449463,-0.0023975372,0.03717041,0.020751953,-0.000009357929,-0.06842041,-0.045074463,-0.035980225,0.03060913,0.00049352646,-0.0013618469,0.018676758,0.00070238113,-0.015472412,-0.035736084,-0.008995056,0.008773804,0.009635925,0.023330688,-0.027008057,-0.0074501038,-0.0040893555,0.010391235,-0.030014038,-0.04119873,-0.06329346,0.049926758,-0.016952515,-0.015045166,-0.0010814667,0.020309448,-0.0034770966,0.05996704,-0.043273926,-0.035491943,0.017654419,0.033325195,-0.015403748,0.03942871,-0.003692627,-0.008995056,-0.012290955,-0.004722595,0.010276794,-0.027023315,-0.0052871704,0.019729614,0.026519775,-0.029541016,-0.05505371,0.007499695,-0.030639648,0.00042963028,-0.016693115,0.03125,0.03543091,0.010482788,0.018081665,0.030441284,0.030960083,-0.008422852,-0.00983429,0.047332764,0.0023212433,0.0052719116]],"meta":{"api_version":{"version":"1"},"billed_units":{"input_tokens":2}},"response_type":"embeddings_floats"}'
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      Via:
+      - 1.1 google
+      access-control-expose-headers:
+      - X-Debug-Trace-ID
+      cache-control:
+      - no-cache, no-store, no-transform, must-revalidate, private, max-age=0
+      content-type:
+      - application/json
+      date:
+      - Fri, 26 Apr 2024 13:51:32 GMT
+      expires:
+      - Thu, 01 Jan 1970 00:00:00 UTC
+      num_chars:
+      - '12'
+      num_tokens:
+      - '2'
+      pragma:
+      - no-cache
+      server:
+      - envoy
+      transfer-encoding:
+      - chunked
+      vary:
+      - Origin
+      x-accel-expires:
+      - '0'
+      x-debug-trace-id:
+      - a4a9481da87382f77e640a8707065fbe
+      x-endpoint-monthly-call-limit:
+      - '1000'
+      x-envoy-upstream-service-time:
+      - '69'
+      x-trial-endpoint-call-limit:
+      - '40'
+      x-trial-endpoint-call-remaining:
+      - '39'
+    status:
+      code: 200
+      message: OK
+version: 1

tests/cohere/cassettes/test_cohere_rerank.yaml

@@ -0,0 +1,76 @@
+interactions:
+- request:
+    body: '{"query": "What is the capital of the United States?", "documents": ["Carson
+      City is the capital city of the American state of Nevada.", "The Commonwealth
+      of the Northern Mariana Islands is a group of islands in the Pacific Ocean.
+      Its capital is Saipan.", "Washington, D.C. (also known as simply Washington
+      or D.C., and officially as the District of Columbia) is the capital of the United
+      States. It is a federal district.", "Capital punishment (the death penalty)
+      has existed in the United States since beforethe United States was a country.
+      As of 2017, capital punishment is legal in 30 of the 50 states."], "model":
+      "rerank-english-v2.0", "top_n": 3}'
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '653'
+      content-type:
+      - application/json
+      host:
+      - api.cohere.ai
+      user-agent:
+      - python-httpx/0.27.0
+      x-fern-language:
+      - Python
+      x-fern-sdk-name:
+      - cohere
+      x-fern-sdk-version:
+      - 5.3.2
+    method: POST
+    uri: https://api.cohere.ai/v1/rerank
+  response:
+    body:
+      string: '{"id":"81824c15-b198-48cb-843c-08058b40a39e","results":[{"index":2,"relevance_score":0.98005307},{"index":3,"relevance_score":0.27904198},{"index":0,"relevance_score":0.10194652}],"meta":{"api_version":{"version":"1"},"billed_units":{"search_units":1}}}'
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      Via:
+      - 1.1 google
+      access-control-expose-headers:
+      - X-Debug-Trace-ID
+      cache-control:
+      - no-cache, no-store, no-transform, must-revalidate, private, max-age=0
+      content-length:
+      - '253'
+      content-type:
+      - application/json
+      date:
+      - Fri, 26 Apr 2024 13:51:19 GMT
+      expires:
+      - Thu, 01 Jan 1970 00:00:00 UTC
+      pragma:
+      - no-cache
+      server:
+      - envoy
+      vary:
+      - Origin
+      x-accel-expires:
+      - '0'
+      x-debug-trace-id:
+      - 8063c83d2452d0fa74f6798659d96da8
+      x-endpoint-monthly-call-limit:
+      - '1000'
+      x-envoy-upstream-service-time:
+      - '40'
+      x-trial-endpoint-call-limit:
+      - '40'
+      x-trial-endpoint-call-remaining:
+      - '39'
+    status:
+      code: 200
+      message: OK
+version: 1

tests/cohere/conftest.py

@@ -0,0 +1,30 @@
+"""Unit tests configuration module."""
+
+import pytest
+import os
+import cohere
+
+from langtrace_python_sdk.instrumentation.cohere.instrumentation import (
+    CohereInstrumentation,
+)
+
+
+@pytest.fixture(autouse=True)
+def environment():
+    if not os.getenv("CO_API_KEY"):
+        os.environ["CO_API_KEY"] = "test_co_api_key"
+
+
+@pytest.fixture
+def cohere_client():
+    return cohere.Client(os.environ["CO_API_KEY"])
+
+
+@pytest.fixture(scope="module")
+def vcr_config():
+    return {"filter_headers": ["authorization"]}
+
+
+@pytest.fixture(scope="session", autouse=True)
+def instrument():
+    CohereInstrumentation().instrument()

tests/cohere/test_cohere_chat.py

@@ -0,0 +1,116 @@
+import json
+from langtrace_python_sdk.constants.instrumentation.cohere import APIS
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
+import pytest
+import importlib
+from tests.utils import assert_response_format, assert_token_count
+
+
+@pytest.mark.vcr
+def test_cohere_chat(cohere_client, exporter):
+    llm_model_value = "command"
+    messages_value = "Tell me a story in 3 sentences or less?"
+    preamble_value = "answer like a pirate"
+    connectors = [{"id": "web-search"}]
+    kwargs = {
+        "model": llm_model_value,
+        "message": messages_value,
+        "preamble": preamble_value,
+        "connectors": connectors,
+        "temperature": 0.1,
+        "chat_history": [
+            {"role": "USER", "message": "Who discovered gravity?"},
+            {
+                "role": "CHATBOT",
+                "message": "The man who is widely credited with discovering gravity is Sir Isaac Newton",
+            },
+        ],
+    }
+    res = cohere_client.chat(**kwargs)
+    spans = exporter.get_finished_spans()
+    cohere_span = spans[-1]
+    assert cohere_span.name == APIS["CHAT_CREATE"]["METHOD"]
+    attributes = cohere_span.attributes
+
+    assert attributes.get("langtrace.sdk.name") == "langtrace-python-sdk"
+    assert attributes.get("langtrace.service.name") == SERVICE_PROVIDERS["COHERE"]
+    assert attributes.get("langtrace.service.type") == "llm"
+    assert attributes.get("langtrace.service.version") == importlib.metadata.version(
+        "cohere"
+    )
+
+    assert attributes.get("langtrace.version") == "1.0.0"
+    assert attributes.get("url.full") == APIS["CHAT_CREATE"]["URL"]
+    assert attributes.get("llm.api") == APIS["CHAT_CREATE"]["ENDPOINT"]
+    assert attributes.get("llm.model") == llm_model_value
+    assert attributes.get("llm.generation_id") == res.generation_id
+    assert attributes.get("llm.temperature") == kwargs.get("temperature")
+    assert attributes.get("llm.stream") is False
+
+    assert json.loads(attributes.get("llm.connectors")) == connectors
+    assert json.loads(attributes.get("llm.prompts"))[-1]["content"] == messages_value
+    assert json.loads(attributes.get("llm.responses"))[-1]["content"] == res.text
+
+    assert_token_count(attributes)
+    assert_response_format(attributes)
+
+
+@pytest.mark.vcr
+def test_cohere_chat_streaming(cohere_client, exporter):
+    llm_model_value = "command"
+    messages_value = "Tell me a story in 3 sentences or less?"
+    preamble_value = "answer like a pirate"
+    connectors = [{"id": "web-search"}]
+    kwargs = {
+        "model": llm_model_value,
+        "message": messages_value,
+        "preamble": preamble_value,
+        "connectors": connectors,
+        "temperature": 0.1,
+        "chat_history": [
+            {"role": "USER", "message": "Who discovered gravity?"},
+            {
+                "role": "CHATBOT",
+                "message": "The man who is widely credited with discovering gravity is Sir Isaac Newton",
+            },
+        ],
+    }
+
+    res = cohere_client.chat_stream(**kwargs)
+
+    chunks_count = 0
+    streamed_response = ""
+    for chunk in res:
+        if chunk.event_type == "text-generation":
+            streamed_response += chunk.text
+        chunks_count += 1
+
+    spans = exporter.get_finished_spans()
+    cohere_span = spans[-1]
+    assert cohere_span.name == APIS["CHAT_STREAM"]["METHOD"]
+    attributes = cohere_span.attributes
+
+    assert attributes.get("langtrace.sdk.name") == "langtrace-python-sdk"
+    assert attributes.get("langtrace.service.name") == SERVICE_PROVIDERS["COHERE"]
+    assert attributes.get("langtrace.service.type") == "llm"
+    assert attributes.get("langtrace.service.version") == importlib.metadata.version(
+        "cohere"
+    )
+
+    assert attributes.get("langtrace.version") == "1.0.0"
+    assert attributes.get("url.full") == APIS["CHAT_STREAM"]["URL"]
+    assert attributes.get("llm.api") == APIS["CHAT_STREAM"]["ENDPOINT"]
+    assert attributes.get("llm.model") == llm_model_value
+    assert attributes.get("llm.temperature") == kwargs.get("temperature")
+    assert attributes.get("llm.stream") is True
+    assert json.loads(attributes.get("llm.connectors")) == connectors
+    assert json.loads(attributes.get("llm.prompts"))[-1]["content"] == messages_value
+    events = cohere_span.events
+    assert events[-1].name == "stream.end"
+    assert len(events) - 2 == chunks_count
+    assert (
+        json.loads(attributes.get("llm.responses"))[-1]["content"] == streamed_response
+    )
+
+    assert_token_count(attributes)
+    assert_response_format(attributes)

tests/cohere/test_cohere_embed.py

@@ -0,0 +1,35 @@
+from langtrace_python_sdk.constants.instrumentation.cohere import APIS
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
+import pytest
+import importlib
+
+
+@pytest.mark.vcr
+def test_cohere_embed(cohere_client, exporter):
+    llm_model_value = "embed-english-v3.0"
+    texts = ["hello", "goodbye"]
+    input_type = "classification"
+
+    kwargs = {
+        "model": llm_model_value,
+        "texts": texts,
+        "input_type": input_type,
+    }
+
+    cohere_client.embed(**kwargs)
+    spans = exporter.get_finished_spans()
+    cohere_span = spans[-1]
+    assert cohere_span.name == APIS["EMBED"]["METHOD"]
+    attributes = cohere_span.attributes
+
+    assert attributes.get("langtrace.sdk.name") == "langtrace-python-sdk"
+    assert attributes.get("langtrace.service.name") == SERVICE_PROVIDERS["COHERE"]
+    assert attributes.get("langtrace.service.type") == "llm"
+    assert attributes.get("langtrace.service.version") == importlib.metadata.version(
+        "cohere"
+    )
+
+    assert attributes.get("langtrace.version") == "1.0.0"
+    assert attributes.get("url.full") == APIS["EMBED"]["URL"]
+    assert attributes.get("llm.api") == APIS["EMBED"]["ENDPOINT"]
+    assert attributes.get("llm.model") == llm_model_value

tests/cohere/test_cohere_rerank.py

@@ -0,0 +1,52 @@
+from langtrace_python_sdk.constants.instrumentation.common import SERVICE_PROVIDERS
+import pytest
+import json
+import importlib
+from langtrace_python_sdk.constants.instrumentation.cohere import APIS
+from tests.utils import assert_response_format, assert_token_count
+import pprint
+
+
+@pytest.mark.vcr
+def test_cohere_rerank(cohere_client, exporter):
+    llm_model_value = "rerank-english-v2.0"
+    query = "What is the capital of the United States?"
+    docs = [
+        "Carson City is the capital city of the American state of Nevada.",
+        "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
+        "Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) is the capital of the United States. It is a federal district.",
+        "Capital punishment (the death penalty) has existed in the United States since beforethe United States was a country. As of 2017, capital punishment is legal in 30 of the 50 states.",
+    ]
+
+    kwargs = {
+        "model": llm_model_value,
+        "query": query,
+        "documents": docs,
+        "top_n": 3,
+    }
+
+    results = cohere_client.rerank(**kwargs)
+    spans = exporter.get_finished_spans()
+    cohere_span = spans[-1]
+    assert cohere_span.name == APIS["RERANK"]["METHOD"]
+    attributes = cohere_span.attributes
+
+    assert attributes.get("langtrace.sdk.name") == "langtrace-python-sdk"
+    assert attributes.get("langtrace.service.name") == SERVICE_PROVIDERS["COHERE"]
+    assert attributes.get("langtrace.service.type") == "llm"
+    assert attributes.get("langtrace.service.version") == importlib.metadata.version(
+        "cohere"
+    )
+
+    assert attributes.get("langtrace.version") == "1.0.0"
+    assert attributes.get("url.full") == APIS["RERANK"]["URL"]
+    assert attributes.get("llm.api") == APIS["RERANK"]["ENDPOINT"]
+    assert attributes.get("llm.model") == llm_model_value
+
+    langtrace_results = json.loads(attributes.get("llm.retrieval.results"))
+    for idx, res in enumerate(results.results):
+        lang_res = json.loads(langtrace_results[idx])
+        assert lang_res["index"] == res.index
+        assert lang_res["relevance_score"] == res.relevance_score
+
+    assert_token_count(attributes)

tests/conftest.py

@@ -15,3 +15,8 @@ def exporter():
     trace.set_tracer_provider(provider)
 
     return exporter
+
+
+@pytest.fixture(autouse=True)
+def clear_exporter(exporter):
+    exporter.clear()

tests/openai/conftest.py

@@ -32,6 +32,3 @@ def vcr_config():
 @pytest.fixture(scope="session", autouse=True)
 def instrument():
     OpenAIInstrumentation().instrument()
-
-    yield
-    # exporter.shutdown()

tests/openai/test_chat_completion.py

@@ -2,6 +2,7 @@ import pytest
 import importlib
 import json
 from langtrace_python_sdk.constants.instrumentation.openai import APIS
+from tests.utils import assert_response_format, assert_token_count
 
 
 @pytest.mark.vcr()
@@ -34,13 +35,15 @@ def test_chat_completion(exporter, openai_client):
     assert attributes.get("llm.prompts") == json.dumps(messages_value)
     assert attributes.get("llm.stream") is False
 
-    tokens = json.loads(attributes.get("llm.token.counts"))
-    output_tokens = tokens.get("output_tokens")
-    prompt_tokens = tokens.get("input_tokens")
-    total_tokens = tokens.get("total_tokens")
+    assert_token_count(attributes)
+    assert_response_format(attributes)
 
-    assert output_tokens and prompt_tokens and total_tokens
-    assert output_tokens + prompt_tokens == total_tokens
+    langtrace_responses = json.loads(attributes.get("llm.responses"))
+    assert isinstance(langtrace_responses, list)
+    for langtrace_response in langtrace_responses:
+        assert isinstance(langtrace_response, dict)
+        assert "role" in langtrace_response
+        assert "content" in langtrace_response
 
     langtrace_responses = json.loads(attributes.get("llm.responses"))
     assert isinstance(langtrace_responses, list)
@@ -88,15 +91,15 @@ def test_chat_completion_streaming(exporter, openai_client):
     events = streaming_span.events
     assert len(events) - 2 == chunk_count  # -2 for start and end events
 
-    # check token usage attributes for stream
-    tokens = json.loads(attributes.get("llm.token.counts"))
+    assert_token_count(attributes)
+    assert_response_format(attributes)
 
-    output_tokens = tokens.get("output_tokens")
-    prompt_tokens = tokens.get("input_tokens")
-    total_tokens = tokens.get("total_tokens")
-
-    assert output_tokens and prompt_tokens and total_tokens
-    assert output_tokens + prompt_tokens == total_tokens
+    langtrace_responses = json.loads(attributes.get("llm.responses"))
+    assert isinstance(langtrace_responses, list)
+    for langtrace_response in langtrace_responses:
+        assert isinstance(langtrace_response, dict)
+        assert "role" in langtrace_response
+        assert "content" in langtrace_response
 
     langtrace_responses = json.loads(attributes.get("llm.responses"))
     assert isinstance(langtrace_responses, list)
@@ -145,19 +148,5 @@ async def test_async_chat_completion_streaming(exporter, async_openai_client):
     events = streaming_span.events
     assert len(events) - 2 == chunk_count  # -2 for start and end events
 
-    # check token usage attributes for stream
-    tokens = json.loads(attributes.get("llm.token.counts"))
-
-    output_tokens = tokens.get("output_tokens")
-    prompt_tokens = tokens.get("input_tokens")
-    total_tokens = tokens.get("total_tokens")
-
-    assert output_tokens and prompt_tokens and total_tokens
-    assert output_tokens + prompt_tokens == total_tokens
-
-    langtrace_responses = json.loads(attributes.get("llm.responses"))
-    assert isinstance(langtrace_responses, list)
-    for langtrace_response in langtrace_responses:
-        assert isinstance(langtrace_response, dict)
-        assert "role" in langtrace_response
-        assert "content" in langtrace_response
+    assert_token_count(attributes)
+    assert_response_format(attributes)

tests/utils.py

@@ -1,6 +1,7 @@
 from unittest.mock import MagicMock, patch
 import json
-   
+
+
 def common_setup(data, method_to_mock=None):
     if method_to_mock:
         service_mock = patch(method_to_mock)
@@ -9,7 +10,6 @@ def common_setup(data, method_to_mock=None):
     else:
         service_mock = MagicMock()
         service_mock.return_value = MagicMock(**data)
-       
 
     tracer = MagicMock()
     span = MagicMock()
@@ -18,4 +18,27 @@ def common_setup(data, method_to_mock=None):
     context_manager_mock.__enter__.return_value = span
     tracer.start_as_current_span.return_value = context_manager_mock
 
-    return service_mock, tracer, span
+    return service_mock, tracer, span
+
+
+def assert_token_count(attributes):
+    tokens = json.loads(attributes.get("llm.token.counts"))
+    output_tokens = tokens.get("output_tokens")
+    prompt_tokens = tokens.get("input_tokens")
+    total_tokens = tokens.get("total_tokens")
+
+    assert (
+        output_tokens is not None
+        and prompt_tokens is not None
+        and total_tokens is not None
+    )
+    assert output_tokens + prompt_tokens == total_tokens
+
+
+def assert_response_format(attributes):
+    langtrace_responses = json.loads(attributes.get("llm.responses"))
+    assert isinstance(langtrace_responses, list)
+    for langtrace_response in langtrace_responses:
+        assert isinstance(langtrace_response, dict)
+        assert "role" in langtrace_response
+        assert "content" in langtrace_response