diff --git a/js/src/tests/evaluate.int.test.ts b/js/src/tests/evaluate.int.test.ts index 4d68b920c..1d62dd6d9 100644 --- a/js/src/tests/evaluate.int.test.ts +++ b/js/src/tests/evaluate.int.test.ts @@ -625,7 +625,7 @@ test("max concurrency works with summary evaluators", async () => { expect(receivedCommentStrings).toEqual(expectedCommentString); }); -test("Target func can be a runnable", async () => { +test.skip("Target func can be a runnable", async () => { const targetFunc = RunnableSequence.from([ RunnableLambda.from((input: Record) => ({ foo: input.input + 1, diff --git a/python/langsmith/schemas.py b/python/langsmith/schemas.py index d859418e7..8ad12b3d0 100644 --- a/python/langsmith/schemas.py +++ b/python/langsmith/schemas.py @@ -17,7 +17,7 @@ ) from uuid import UUID -from typing_extensions import TypedDict +from typing_extensions import NotRequired, TypedDict try: from pydantic.v1 import ( # type: ignore[import] @@ -891,3 +891,64 @@ class PromptSortField(str, Enum): """Last updated time.""" num_likes = "num_likes" """Number of likes.""" + + +class InputTokenDetails(TypedDict, total=False): + """Breakdown of input token counts. + + Does *not* need to sum to full input token count. Does *not* need to have all keys. + """ + + audio: int + """Audio input tokens.""" + cache_creation: int + """Input tokens that were cached and there was a cache miss. + + Since there was a cache miss, the cache was created from these tokens. + """ + cache_read: int + """Input tokens that were cached and there was a cache hit. + + Since there was a cache hit, the tokens were read from the cache. More precisely, + the model state given these tokens was read from the cache. + """ + + +class OutputTokenDetails(TypedDict, total=False): + """Breakdown of output token counts. + + Does *not* need to sum to full output token count. Does *not* need to have all keys. + """ + + audio: int + """Audio output tokens.""" + reasoning: int + """Reasoning output tokens. + + Tokens generated by the model in a chain of thought process (i.e. by OpenAI's o1 + models) that are not returned as part of model output. + """ + + +class UsageMetadata(TypedDict): + """Usage metadata for a message, such as token counts. + + This is a standard representation of token usage that is consistent across models. + """ + + input_tokens: int + """Count of input (or prompt) tokens. Sum of all input token types.""" + output_tokens: int + """Count of output (or completion) tokens. Sum of all output token types.""" + total_tokens: int + """Total token count. Sum of input_tokens + output_tokens.""" + input_token_details: NotRequired[InputTokenDetails] + """Breakdown of input token counts. + + Does *not* need to sum to full input token count. Does *not* need to have all keys. + """ + output_token_details: NotRequired[OutputTokenDetails] + """Breakdown of output token counts. + + Does *not* need to sum to full output token count. Does *not* need to have all keys. + """ diff --git a/python/langsmith/wrappers/_openai.py b/python/langsmith/wrappers/_openai.py index 014d364cd..799a64493 100644 --- a/python/langsmith/wrappers/_openai.py +++ b/python/langsmith/wrappers/_openai.py @@ -21,6 +21,7 @@ from langsmith import client as ls_client from langsmith import run_helpers +from langsmith.schemas import InputTokenDetails, OutputTokenDetails, UsageMetadata if TYPE_CHECKING: from openai import AsyncOpenAI, OpenAI @@ -141,6 +142,12 @@ def _reduce_chat(all_chunks: List[ChatCompletionChunk]) -> dict: ] else: d = {"choices": [{"message": {"role": "assistant", "content": ""}}]} + # streamed outputs don't go through `process_outputs` + # so we need to flatten metadata here + oai_token_usage = d.pop("usage", None) + d["usage_metadata"] = ( + _create_usage_metadata(oai_token_usage) if oai_token_usage else None + ) return d @@ -160,12 +167,59 @@ def _reduce_completions(all_chunks: List[Completion]) -> dict: return d +def _create_usage_metadata(oai_token_usage: dict) -> UsageMetadata: + input_tokens = oai_token_usage.get("prompt_tokens") or 0 + output_tokens = oai_token_usage.get("completion_tokens") or 0 + total_tokens = oai_token_usage.get("total_tokens") or input_tokens + output_tokens + input_token_details: dict = { + "audio": (oai_token_usage.get("prompt_tokens_details") or {}).get( + "audio_tokens" + ), + "cache_read": (oai_token_usage.get("prompt_tokens_details") or {}).get( + "cached_tokens" + ), + } + output_token_details: dict = { + "audio": (oai_token_usage.get("completion_tokens_details") or {}).get( + "audio_tokens" + ), + "reasoning": (oai_token_usage.get("completion_tokens_details") or {}).get( + "reasoning_tokens" + ), + } + return UsageMetadata( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens, + input_token_details=InputTokenDetails( + **{k: v for k, v in input_token_details.items() if v is not None} + ), + output_token_details=OutputTokenDetails( + **{k: v for k, v in output_token_details.items() if v is not None} + ), + ) + + +def _process_chat_completion(outputs: Any): + try: + rdict = outputs.model_dump() + oai_token_usage = rdict.pop("usage", None) + rdict["usage_metadata"] = ( + _create_usage_metadata(oai_token_usage) if oai_token_usage else None + ) + return rdict + except BaseException as e: + logger.debug(f"Error processing chat completion: {e}") + return {"output": outputs} + + def _get_wrapper( original_create: Callable, name: str, reduce_fn: Callable, tracing_extra: Optional[TracingExtra] = None, invocation_params_fn: Optional[Callable] = None, + process_outputs: Optional[Callable] = None, ) -> Callable: textra = tracing_extra or {} @@ -177,6 +231,7 @@ def create(*args, stream: bool = False, **kwargs): reduce_fn=reduce_fn if stream else None, process_inputs=_strip_not_given, _invocation_params_fn=invocation_params_fn, + process_outputs=process_outputs, **textra, ) @@ -191,6 +246,7 @@ async def acreate(*args, stream: bool = False, **kwargs): reduce_fn=reduce_fn if stream else None, process_inputs=_strip_not_given, _invocation_params_fn=invocation_params_fn, + process_outputs=process_outputs, **textra, ) return await decorator(original_create)(*args, stream=stream, **kwargs) @@ -232,6 +288,7 @@ def wrap_openai( _reduce_chat, tracing_extra=tracing_extra, invocation_params_fn=functools.partial(_infer_invocation_params, "chat"), + process_outputs=_process_chat_completion, ) client.completions.create = _get_wrapper( # type: ignore[method-assign] client.completions.create, diff --git a/python/tests/integration_tests/test_data/langsmith_py_wrap_openai_.json b/python/tests/integration_tests/test_data/langsmith_py_wrap_openai_.json new file mode 100644 index 000000000..176b4b94c --- /dev/null +++ b/python/tests/integration_tests/test_data/langsmith_py_wrap_openai_.json @@ -0,0 +1,120 @@ +{ + "post": [ + { + "id": "d0d84d31-923d-4cb5-94a8-40a0a0087578", + "start_time": "2024-10-11T20:58:23.298773+00:00", + "extra": { + "metadata": { + "ls_method": "traceable", + "ls_provider": "openai", + "ls_model_type": "chat", + "ls_model_name": "gpt-4o-mini", + "revision_id": "v0.1.82-381-g03d9e1a-dirty" + }, + "runtime": { + "sdk": "langsmith-py", + "sdk_version": "0.1.131", + "library": "langsmith", + "platform": "macOS-13.2-arm64-arm-64bit", + "runtime": "python", + "py_implementation": "CPython", + "runtime_version": "3.11.7", + "langchain_version": "0.2.9", + "langchain_core_version": "0.2.21" + } + }, + "serialized": { + "name": "ChatOpenAI", + "signature": "(*, messages: 'Iterable[ChatCompletionMessageParam]', model: 'Union[str, ChatModel]', frequency_penalty: 'Optional[float] | NotGiven' = NOT_GIVEN, function_call: 'completion_create_params.FunctionCall | NotGiven' = NOT_GIVEN, functions: 'Iterable[completion_create_params.Function] | NotGiven' = NOT_GIVEN, logit_bias: 'Optional[Dict[str, int]] | NotGiven' = NOT_GIVEN, logprobs: 'Optional[bool] | NotGiven' = NOT_GIVEN, max_completion_tokens: 'Optional[int] | NotGiven' = NOT_GIVEN, max_tokens: 'Optional[int] | NotGiven' = NOT_GIVEN, n: 'Optional[int] | NotGiven' = NOT_GIVEN, parallel_tool_calls: 'bool | NotGiven' = NOT_GIVEN, presence_penalty: 'Optional[float] | NotGiven' = NOT_GIVEN, response_format: 'completion_create_params.ResponseFormat | NotGiven' = NOT_GIVEN, seed: 'Optional[int] | NotGiven' = NOT_GIVEN, service_tier: \"Optional[Literal['auto', 'default']] | NotGiven\" = NOT_GIVEN, stop: 'Union[Optional[str], List[str]] | NotGiven' = NOT_GIVEN, stream: 'Optional[Literal[False]] | Literal[True] | NotGiven' = NOT_GIVEN, stream_options: 'Optional[ChatCompletionStreamOptionsParam] | NotGiven' = NOT_GIVEN, temperature: 'Optional[float] | NotGiven' = NOT_GIVEN, tool_choice: 'ChatCompletionToolChoiceOptionParam | NotGiven' = NOT_GIVEN, tools: 'Iterable[ChatCompletionToolParam] | NotGiven' = NOT_GIVEN, top_logprobs: 'Optional[int] | NotGiven' = NOT_GIVEN, top_p: 'Optional[float] | NotGiven' = NOT_GIVEN, user: 'str | NotGiven' = NOT_GIVEN, extra_headers: 'Headers | None' = None, extra_query: 'Query | None' = None, extra_body: 'Body | None' = None, timeout: 'float | httpx.Timeout | None | NotGiven' = NOT_GIVEN) -> 'ChatCompletion | AsyncStream[ChatCompletionChunk]'", + "doc": null + }, + "events": [], + "tags": [], + "attachments": {}, + "dotted_order": "20241011T205823298773Zd0d84d31-923d-4cb5-94a8-40a0a0087578", + "trace_id": "d0d84d31-923d-4cb5-94a8-40a0a0087578", + "outputs": {}, + "session_name": "default", + "name": "ChatOpenAI", + "inputs": { + "messages": [ + { + "role": "user", + "content": "howdy" + } + ], + "model": "gpt-4o-mini", + "stream": false, + "extra_headers": null, + "extra_query": null, + "extra_body": null + }, + "run_type": "llm" + } + ], + "patch": [ + { + "id": "d0d84d31-923d-4cb5-94a8-40a0a0087578", + "name": "ChatOpenAI", + "trace_id": "d0d84d31-923d-4cb5-94a8-40a0a0087578", + "parent_run_id": null, + "dotted_order": "20241011T205823298773Zd0d84d31-923d-4cb5-94a8-40a0a0087578", + "tags": [], + "extra": { + "metadata": { + "ls_method": "traceable", + "ls_provider": "openai", + "ls_model_type": "chat", + "ls_model_name": "gpt-4o-mini", + "revision_id": "v0.1.82-381-g03d9e1a-dirty" + }, + "runtime": { + "sdk": "langsmith-py", + "sdk_version": "0.1.131", + "library": "langsmith", + "platform": "macOS-13.2-arm64-arm-64bit", + "runtime": "python", + "py_implementation": "CPython", + "runtime_version": "3.11.7", + "langchain_version": "0.2.9", + "langchain_core_version": "0.2.21" + } + }, + "end_time": "2024-10-11T20:58:24.417106+00:00", + "outputs": { + "id": "chatcmpl-AHH0KBvLG7Wq3wfSEGQuxh0xE07Fl", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": null, + "message": { + "content": "Howdy! How can I assist you today?", + "refusal": null, + "role": "assistant", + "function_call": null, + "tool_calls": null + } + } + ], + "created": 1728680304, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion", + "service_tier": null, + "system_fingerprint": "fp_e2bde53e6e", + "usage_metadata": { + "input_tokens": 9, + "output_tokens": 9, + "total_tokens": 18, + "input_token_details": { + "cache_read": 0 + }, + "output_token_details": { + "reasoning": 0 + } + } + }, + "events": [] + } + ] +} \ No newline at end of file diff --git a/python/tests/integration_tests/test_data/langsmith_py_wrap_openai_reasoning.json b/python/tests/integration_tests/test_data/langsmith_py_wrap_openai_reasoning.json new file mode 100644 index 000000000..706e86886 --- /dev/null +++ b/python/tests/integration_tests/test_data/langsmith_py_wrap_openai_reasoning.json @@ -0,0 +1,120 @@ +{ + "post": [ + { + "id": "a8b34ded-ccd2-4fb7-bccb-9cd625066a14", + "start_time": "2024-10-11T20:58:24.544431+00:00", + "extra": { + "metadata": { + "ls_method": "traceable", + "ls_provider": "openai", + "ls_model_type": "chat", + "ls_model_name": "o1-mini", + "revision_id": "v0.1.82-381-g03d9e1a-dirty" + }, + "runtime": { + "sdk": "langsmith-py", + "sdk_version": "0.1.131", + "library": "langsmith", + "platform": "macOS-13.2-arm64-arm-64bit", + "runtime": "python", + "py_implementation": "CPython", + "runtime_version": "3.11.7", + "langchain_version": "0.2.9", + "langchain_core_version": "0.2.21" + } + }, + "serialized": { + "name": "ChatOpenAI", + "signature": "(*, messages: 'Iterable[ChatCompletionMessageParam]', model: 'Union[str, ChatModel]', frequency_penalty: 'Optional[float] | NotGiven' = NOT_GIVEN, function_call: 'completion_create_params.FunctionCall | NotGiven' = NOT_GIVEN, functions: 'Iterable[completion_create_params.Function] | NotGiven' = NOT_GIVEN, logit_bias: 'Optional[Dict[str, int]] | NotGiven' = NOT_GIVEN, logprobs: 'Optional[bool] | NotGiven' = NOT_GIVEN, max_completion_tokens: 'Optional[int] | NotGiven' = NOT_GIVEN, max_tokens: 'Optional[int] | NotGiven' = NOT_GIVEN, n: 'Optional[int] | NotGiven' = NOT_GIVEN, parallel_tool_calls: 'bool | NotGiven' = NOT_GIVEN, presence_penalty: 'Optional[float] | NotGiven' = NOT_GIVEN, response_format: 'completion_create_params.ResponseFormat | NotGiven' = NOT_GIVEN, seed: 'Optional[int] | NotGiven' = NOT_GIVEN, service_tier: \"Optional[Literal['auto', 'default']] | NotGiven\" = NOT_GIVEN, stop: 'Union[Optional[str], List[str]] | NotGiven' = NOT_GIVEN, stream: 'Optional[Literal[False]] | Literal[True] | NotGiven' = NOT_GIVEN, stream_options: 'Optional[ChatCompletionStreamOptionsParam] | NotGiven' = NOT_GIVEN, temperature: 'Optional[float] | NotGiven' = NOT_GIVEN, tool_choice: 'ChatCompletionToolChoiceOptionParam | NotGiven' = NOT_GIVEN, tools: 'Iterable[ChatCompletionToolParam] | NotGiven' = NOT_GIVEN, top_logprobs: 'Optional[int] | NotGiven' = NOT_GIVEN, top_p: 'Optional[float] | NotGiven' = NOT_GIVEN, user: 'str | NotGiven' = NOT_GIVEN, extra_headers: 'Headers | None' = None, extra_query: 'Query | None' = None, extra_body: 'Body | None' = None, timeout: 'float | httpx.Timeout | None | NotGiven' = NOT_GIVEN) -> 'ChatCompletion | AsyncStream[ChatCompletionChunk]'", + "doc": null + }, + "events": [], + "tags": [], + "attachments": {}, + "dotted_order": "20241011T205824544431Za8b34ded-ccd2-4fb7-bccb-9cd625066a14", + "trace_id": "a8b34ded-ccd2-4fb7-bccb-9cd625066a14", + "outputs": {}, + "session_name": "default", + "name": "ChatOpenAI", + "inputs": { + "messages": [ + { + "role": "user", + "content": "Write a bash script that takes a matrix represented as a string with format '[1,2],[3,4],[5,6]' and prints the transpose in the same format." + } + ], + "model": "o1-mini", + "stream": false, + "extra_headers": null, + "extra_query": null, + "extra_body": null + }, + "run_type": "llm" + } + ], + "patch": [ + { + "id": "a8b34ded-ccd2-4fb7-bccb-9cd625066a14", + "name": "ChatOpenAI", + "trace_id": "a8b34ded-ccd2-4fb7-bccb-9cd625066a14", + "parent_run_id": null, + "dotted_order": "20241011T205824544431Za8b34ded-ccd2-4fb7-bccb-9cd625066a14", + "tags": [], + "extra": { + "metadata": { + "ls_method": "traceable", + "ls_provider": "openai", + "ls_model_type": "chat", + "ls_model_name": "o1-mini", + "revision_id": "v0.1.82-381-g03d9e1a-dirty" + }, + "runtime": { + "sdk": "langsmith-py", + "sdk_version": "0.1.131", + "library": "langsmith", + "platform": "macOS-13.2-arm64-arm-64bit", + "runtime": "python", + "py_implementation": "CPython", + "runtime_version": "3.11.7", + "langchain_version": "0.2.9", + "langchain_core_version": "0.2.21" + } + }, + "end_time": "2024-10-11T20:58:39.682524+00:00", + "outputs": { + "id": "chatcmpl-AHH0LWUyAupsCrDZu564ZHwRbNQeZ", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": null, + "message": { + "content": "Certainly! Below is a **Bash script** that takes a matrix represented as a string in the format `\"[1,2],[3,4],[5,6]\"` and prints its transpose in the same format. The script uses `awk` to handle the parsing and transposition logic efficiently.\n\n### **Script: `transpose_matrix.sh`**\n\n```bash\n#!/bin/bash\n\n# Check if exactly one argument is provided\nif [ \"$#\" -ne 1 ]; then\n echo \"Usage: $0 '[1,2],[3,4],[5,6]'\"\n exit 1\nfi\n\ninput=\"$1\"\n\n# Use awk to parse the input and perform the transpose\necho \"$input\" | awk '\nBEGIN {\n # Define the field separator to split the input into rows\n FS=\"\\\\],\\\\[|\\\\[|\\\\]\"\n}\n\n{\n row = 0\n # Iterate over each field (row)\n for (i = 1; i <= NF; i++) {\n if ($i != \"\") {\n row++\n # Split the row into individual elements based on comma\n split($i, elements, \",\")\n for (j = 1; j <= length(elements); j++) {\n # Store elements in a 2D array\n matrix[j, row] = elements[j]\n # Keep track of the maximum number of columns and rows\n if (j > max_col) max_col = j\n if (row > max_row) max_row = row\n }\n }\n }\n}\n\nEND {\n # Initialize an empty string to build the output\n output = \"\"\n # Iterate over each column to create transposed rows\n for (i = 1; i <= max_col; i++) {\n output = output \"[\"\n for (j = 1; j <= max_row; j++) {\n output = output matrix[i, j]\n if (j < max_row) {\n output = output \",\"\n }\n }\n output = output \"]\"\n if (i < max_col) {\n output = output \",\"\n }\n # Append the transposed row to the final output\n transposed = transposed output\n }\n # Print the final transposed matrix\n print transposed\n}\n'\n```\n\n### **How It Works**\n\n1. **Input Validation:**\n - The script first checks if exactly one argument is provided. If not, it displays usage instructions and exits.\n\n2. **Parsing with `awk`:**\n - **Field Separator (`FS`):**\n - The `FS` is set to handle the input format by splitting the string into individual rows. It looks for `\"],[\"`, `\"[\"`, or `\"]\"` as separators.\n \n - **Reading Rows and Columns:**\n - For each row, the script splits the elements by commas and stores them in a 2D array `matrix[j, row]`, where `j` is the column index and `row` is the row index.\n - It also keeps track of the maximum number of columns (`max_col`) and rows (`max_row`) to handle matrices of varying sizes.\n \n - **Transposing the Matrix:**\n - In the `END` block, the script iterates over each column and constructs transposed rows by collecting elements from each original row.\n - It formats the output to match the input style, enclosing each transposed row in square brackets and separating them with commas.\n\n3. **Execution:**\n - Make the script executable:\n ```bash\n chmod +x transpose_matrix.sh\n ```\n - Run the script with a matrix string as an argument:\n ```bash\n ./transpose_matrix.sh \"[1,2],[3,4],[5,6]\"\n ```\n - **Output:**\n ```\n [1,3,5],[2,4,6]\n ```\n\n### **Examples**\n\n1. **Square Matrix:**\n ```bash\n ./transpose_matrix.sh \"[1,2],[3,4]\"\n ```\n **Output:**\n ```\n [1,3],[2,4]\n ```\n\n2. **Non-Square Matrix:**\n ```bash\n ./transpose_matrix.sh \"[1,2,3],[4,5,6]\"\n ```\n **Output:**\n ```\n [1,4],[2,5],[3,6]\n ```\n\n3. **Matrix with Negative Numbers and Multiple Digits:**\n ```bash\n ./transpose_matrix.sh \"[10,-2,33],[4,5,-6]\"\n ```\n **Output:**\n ```\n [10,4],[-2,5],[33,-6]\n ```\n\n### **Notes**\n\n- **Robustness:**\n - The script assumes that the input is well-formed, with each row enclosed in square brackets and elements separated by commas.\n - It can handle matrices that are not square (i.e., different numbers of rows and columns).\n\n- **Dependencies:**\n - The script relies on `awk`, which is commonly available in Unix-like environments.\n\nFeel free to modify and enhance the script based on your specific needs!", + "refusal": null, + "role": "assistant", + "function_call": null, + "tool_calls": null + } + } + ], + "created": 1728680305, + "model": "o1-mini-2024-09-12", + "object": "chat.completion", + "service_tier": null, + "system_fingerprint": "fp_692002f015", + "usage_metadata": { + "input_tokens": 43, + "output_tokens": 2497, + "total_tokens": 2540, + "input_token_details": { + "cache_read": 0 + }, + "output_token_details": { + "reasoning": 1408 + } + } + }, + "events": [] + } + ] +} \ No newline at end of file diff --git a/python/tests/integration_tests/test_data/langsmith_py_wrap_openai_stream.json b/python/tests/integration_tests/test_data/langsmith_py_wrap_openai_stream.json new file mode 100644 index 000000000..96f165364 --- /dev/null +++ b/python/tests/integration_tests/test_data/langsmith_py_wrap_openai_stream.json @@ -0,0 +1,376 @@ +{ + "post": [ + { + "id": "fe8ffecb-72ce-4cd2-bdb7-01f34654c391", + "start_time": "2024-10-11T20:58:20.695375+00:00", + "extra": { + "metadata": { + "ls_method": "traceable", + "ls_provider": "openai", + "ls_model_type": "chat", + "ls_model_name": "gpt-4o-mini", + "revision_id": "v0.1.82-381-g03d9e1a-dirty" + }, + "runtime": { + "sdk": "langsmith-py", + "sdk_version": "0.1.131", + "library": "langsmith", + "platform": "macOS-13.2-arm64-arm-64bit", + "runtime": "python", + "py_implementation": "CPython", + "runtime_version": "3.11.7", + "langchain_version": "0.2.9", + "langchain_core_version": "0.2.21" + } + }, + "serialized": { + "name": "ChatOpenAI", + "signature": "(*, messages: 'Iterable[ChatCompletionMessageParam]', model: 'Union[str, ChatModel]', frequency_penalty: 'Optional[float] | NotGiven' = NOT_GIVEN, function_call: 'completion_create_params.FunctionCall | NotGiven' = NOT_GIVEN, functions: 'Iterable[completion_create_params.Function] | NotGiven' = NOT_GIVEN, logit_bias: 'Optional[Dict[str, int]] | NotGiven' = NOT_GIVEN, logprobs: 'Optional[bool] | NotGiven' = NOT_GIVEN, max_completion_tokens: 'Optional[int] | NotGiven' = NOT_GIVEN, max_tokens: 'Optional[int] | NotGiven' = NOT_GIVEN, n: 'Optional[int] | NotGiven' = NOT_GIVEN, parallel_tool_calls: 'bool | NotGiven' = NOT_GIVEN, presence_penalty: 'Optional[float] | NotGiven' = NOT_GIVEN, response_format: 'completion_create_params.ResponseFormat | NotGiven' = NOT_GIVEN, seed: 'Optional[int] | NotGiven' = NOT_GIVEN, service_tier: \"Optional[Literal['auto', 'default']] | NotGiven\" = NOT_GIVEN, stop: 'Union[Optional[str], List[str]] | NotGiven' = NOT_GIVEN, stream: 'Optional[Literal[False]] | Literal[True] | NotGiven' = NOT_GIVEN, stream_options: 'Optional[ChatCompletionStreamOptionsParam] | NotGiven' = NOT_GIVEN, temperature: 'Optional[float] | NotGiven' = NOT_GIVEN, tool_choice: 'ChatCompletionToolChoiceOptionParam | NotGiven' = NOT_GIVEN, tools: 'Iterable[ChatCompletionToolParam] | NotGiven' = NOT_GIVEN, top_logprobs: 'Optional[int] | NotGiven' = NOT_GIVEN, top_p: 'Optional[float] | NotGiven' = NOT_GIVEN, user: 'str | NotGiven' = NOT_GIVEN, extra_headers: 'Headers | None' = None, extra_query: 'Query | None' = None, extra_body: 'Body | None' = None, timeout: 'float | httpx.Timeout | None | NotGiven' = NOT_GIVEN) -> 'ChatCompletion | AsyncStream[ChatCompletionChunk]'", + "doc": null + }, + "events": [], + "tags": [], + "attachments": {}, + "dotted_order": "20241011T205820695375Zfe8ffecb-72ce-4cd2-bdb7-01f34654c391", + "trace_id": "fe8ffecb-72ce-4cd2-bdb7-01f34654c391", + "outputs": {}, + "session_name": "default", + "name": "ChatOpenAI", + "inputs": { + "messages": [ + { + "role": "user", + "content": "howdy" + } + ], + "model": "gpt-4o-mini", + "stream": true, + "stream_options": { + "include_usage": true + }, + "extra_headers": null, + "extra_query": null, + "extra_body": null + }, + "run_type": "llm" + } + ], + "patch": [ + { + "id": "fe8ffecb-72ce-4cd2-bdb7-01f34654c391", + "name": "ChatOpenAI", + "trace_id": "fe8ffecb-72ce-4cd2-bdb7-01f34654c391", + "parent_run_id": null, + "dotted_order": "20241011T205820695375Zfe8ffecb-72ce-4cd2-bdb7-01f34654c391", + "tags": [], + "extra": { + "metadata": { + "ls_method": "traceable", + "ls_provider": "openai", + "ls_model_type": "chat", + "ls_model_name": "gpt-4o-mini", + "revision_id": "v0.1.82-381-g03d9e1a-dirty" + }, + "runtime": { + "sdk": "langsmith-py", + "sdk_version": "0.1.131", + "library": "langsmith", + "platform": "macOS-13.2-arm64-arm-64bit", + "runtime": "python", + "py_implementation": "CPython", + "runtime_version": "3.11.7", + "langchain_version": "0.2.9", + "langchain_core_version": "0.2.21" + } + }, + "end_time": "2024-10-11T20:58:22.023816+00:00", + "outputs": { + "id": "chatcmpl-AHH0HKxF2K5Rnu1DJ51k9CPTcerd1", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": { + "role": "assistant", + "content": "Howdy! How can I assist you today?" + } + } + ], + "created": 1728680301, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "service_tier": null, + "system_fingerprint": "fp_8552ec53e1", + "usage_metadata": { + "input_tokens": 9, + "output_tokens": 9, + "total_tokens": 18, + "input_token_details": { + "cache_read": 0 + }, + "output_token_details": { + "reasoning": 0 + } + } + }, + "events": [ + { + "name": "new_token", + "time": "2024-10-11T20:58:21.933794+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0HKxF2K5Rnu1DJ51k9CPTcerd1", + "choices": [ + { + "delta": { + "content": "", + "role": "assistant" + }, + "index": 0 + } + ], + "created": 1728680301, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_8552ec53e1" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:21.934186+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0HKxF2K5Rnu1DJ51k9CPTcerd1", + "choices": [ + { + "delta": { + "content": "Howdy" + }, + "index": 0 + } + ], + "created": 1728680301, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_8552ec53e1" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:21.955034+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0HKxF2K5Rnu1DJ51k9CPTcerd1", + "choices": [ + { + "delta": { + "content": "!" + }, + "index": 0 + } + ], + "created": 1728680301, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_8552ec53e1" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:21.955547+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0HKxF2K5Rnu1DJ51k9CPTcerd1", + "choices": [ + { + "delta": { + "content": " How" + }, + "index": 0 + } + ], + "created": 1728680301, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_8552ec53e1" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:22.005714+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0HKxF2K5Rnu1DJ51k9CPTcerd1", + "choices": [ + { + "delta": { + "content": " can" + }, + "index": 0 + } + ], + "created": 1728680301, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_8552ec53e1" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:22.007009+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0HKxF2K5Rnu1DJ51k9CPTcerd1", + "choices": [ + { + "delta": { + "content": " I" + }, + "index": 0 + } + ], + "created": 1728680301, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_8552ec53e1" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:22.008457+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0HKxF2K5Rnu1DJ51k9CPTcerd1", + "choices": [ + { + "delta": { + "content": " assist" + }, + "index": 0 + } + ], + "created": 1728680301, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_8552ec53e1" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:22.008855+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0HKxF2K5Rnu1DJ51k9CPTcerd1", + "choices": [ + { + "delta": { + "content": " you" + }, + "index": 0 + } + ], + "created": 1728680301, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_8552ec53e1" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:22.010922+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0HKxF2K5Rnu1DJ51k9CPTcerd1", + "choices": [ + { + "delta": { + "content": " today" + }, + "index": 0 + } + ], + "created": 1728680301, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_8552ec53e1" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:22.011337+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0HKxF2K5Rnu1DJ51k9CPTcerd1", + "choices": [ + { + "delta": { + "content": "?" + }, + "index": 0 + } + ], + "created": 1728680301, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_8552ec53e1" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:22.012554+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0HKxF2K5Rnu1DJ51k9CPTcerd1", + "choices": [ + { + "delta": {}, + "finish_reason": "stop", + "index": 0 + } + ], + "created": 1728680301, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_8552ec53e1" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:22.015478+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0HKxF2K5Rnu1DJ51k9CPTcerd1", + "choices": [], + "created": 1728680301, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_8552ec53e1", + "usage": { + "completion_tokens": 9, + "prompt_tokens": 9, + "total_tokens": 18, + "completion_tokens_details": { + "reasoning_tokens": 0 + }, + "prompt_tokens_details": { + "cached_tokens": 0 + } + } + } + } + } + ] + } + ] +} \ No newline at end of file diff --git a/python/tests/integration_tests/test_data/langsmith_py_wrap_openai_stream_no_usage.json b/python/tests/integration_tests/test_data/langsmith_py_wrap_openai_stream_no_usage.json new file mode 100644 index 000000000..150a3b79f --- /dev/null +++ b/python/tests/integration_tests/test_data/langsmith_py_wrap_openai_stream_no_usage.json @@ -0,0 +1,338 @@ +{ + "post": [ + { + "id": "de56b9f0-eed2-4195-8786-c6dc0fa897e3", + "start_time": "2024-10-11T20:58:22.254895+00:00", + "extra": { + "metadata": { + "ls_method": "traceable", + "ls_provider": "openai", + "ls_model_type": "chat", + "ls_model_name": "gpt-4o-mini", + "revision_id": "v0.1.82-381-g03d9e1a-dirty" + }, + "runtime": { + "sdk": "langsmith-py", + "sdk_version": "0.1.131", + "library": "langsmith", + "platform": "macOS-13.2-arm64-arm-64bit", + "runtime": "python", + "py_implementation": "CPython", + "runtime_version": "3.11.7", + "langchain_version": "0.2.9", + "langchain_core_version": "0.2.21" + } + }, + "serialized": { + "name": "ChatOpenAI", + "signature": "(*, messages: 'Iterable[ChatCompletionMessageParam]', model: 'Union[str, ChatModel]', frequency_penalty: 'Optional[float] | NotGiven' = NOT_GIVEN, function_call: 'completion_create_params.FunctionCall | NotGiven' = NOT_GIVEN, functions: 'Iterable[completion_create_params.Function] | NotGiven' = NOT_GIVEN, logit_bias: 'Optional[Dict[str, int]] | NotGiven' = NOT_GIVEN, logprobs: 'Optional[bool] | NotGiven' = NOT_GIVEN, max_completion_tokens: 'Optional[int] | NotGiven' = NOT_GIVEN, max_tokens: 'Optional[int] | NotGiven' = NOT_GIVEN, n: 'Optional[int] | NotGiven' = NOT_GIVEN, parallel_tool_calls: 'bool | NotGiven' = NOT_GIVEN, presence_penalty: 'Optional[float] | NotGiven' = NOT_GIVEN, response_format: 'completion_create_params.ResponseFormat | NotGiven' = NOT_GIVEN, seed: 'Optional[int] | NotGiven' = NOT_GIVEN, service_tier: \"Optional[Literal['auto', 'default']] | NotGiven\" = NOT_GIVEN, stop: 'Union[Optional[str], List[str]] | NotGiven' = NOT_GIVEN, stream: 'Optional[Literal[False]] | Literal[True] | NotGiven' = NOT_GIVEN, stream_options: 'Optional[ChatCompletionStreamOptionsParam] | NotGiven' = NOT_GIVEN, temperature: 'Optional[float] | NotGiven' = NOT_GIVEN, tool_choice: 'ChatCompletionToolChoiceOptionParam | NotGiven' = NOT_GIVEN, tools: 'Iterable[ChatCompletionToolParam] | NotGiven' = NOT_GIVEN, top_logprobs: 'Optional[int] | NotGiven' = NOT_GIVEN, top_p: 'Optional[float] | NotGiven' = NOT_GIVEN, user: 'str | NotGiven' = NOT_GIVEN, extra_headers: 'Headers | None' = None, extra_query: 'Query | None' = None, extra_body: 'Body | None' = None, timeout: 'float | httpx.Timeout | None | NotGiven' = NOT_GIVEN) -> 'ChatCompletion | AsyncStream[ChatCompletionChunk]'", + "doc": null + }, + "events": [], + "tags": [], + "attachments": {}, + "dotted_order": "20241011T205822254895Zde56b9f0-eed2-4195-8786-c6dc0fa897e3", + "trace_id": "de56b9f0-eed2-4195-8786-c6dc0fa897e3", + "outputs": {}, + "session_name": "default", + "name": "ChatOpenAI", + "inputs": { + "messages": [ + { + "role": "user", + "content": "howdy" + } + ], + "model": "gpt-4o-mini", + "stream": true, + "extra_headers": null, + "extra_query": null, + "extra_body": null + }, + "run_type": "llm" + } + ], + "patch": [ + { + "id": "de56b9f0-eed2-4195-8786-c6dc0fa897e3", + "name": "ChatOpenAI", + "trace_id": "de56b9f0-eed2-4195-8786-c6dc0fa897e3", + "parent_run_id": null, + "dotted_order": "20241011T205822254895Zde56b9f0-eed2-4195-8786-c6dc0fa897e3", + "tags": [], + "extra": { + "metadata": { + "ls_method": "traceable", + "ls_provider": "openai", + "ls_model_type": "chat", + "ls_model_name": "gpt-4o-mini", + "revision_id": "v0.1.82-381-g03d9e1a-dirty" + }, + "runtime": { + "sdk": "langsmith-py", + "sdk_version": "0.1.131", + "library": "langsmith", + "platform": "macOS-13.2-arm64-arm-64bit", + "runtime": "python", + "py_implementation": "CPython", + "runtime_version": "3.11.7", + "langchain_version": "0.2.9", + "langchain_core_version": "0.2.21" + } + }, + "end_time": "2024-10-11T20:58:23.181899+00:00", + "outputs": { + "id": "chatcmpl-AHH0Ik2ZQY05uutXjxSaS6C3nvYfy", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": { + "role": "assistant", + "content": "Howdy! How can I assist you today?" + } + } + ], + "created": 1728680302, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "service_tier": null, + "system_fingerprint": "fp_e2bde53e6e", + "usage_metadata": null + }, + "events": [ + { + "name": "new_token", + "time": "2024-10-11T20:58:23.044675+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0Ik2ZQY05uutXjxSaS6C3nvYfy", + "choices": [ + { + "delta": { + "content": "", + "role": "assistant" + }, + "index": 0 + } + ], + "created": 1728680302, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_e2bde53e6e" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:23.045159+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0Ik2ZQY05uutXjxSaS6C3nvYfy", + "choices": [ + { + "delta": { + "content": "Howdy" + }, + "index": 0 + } + ], + "created": 1728680302, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_e2bde53e6e" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:23.076141+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0Ik2ZQY05uutXjxSaS6C3nvYfy", + "choices": [ + { + "delta": { + "content": "!" + }, + "index": 0 + } + ], + "created": 1728680302, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_e2bde53e6e" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:23.076801+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0Ik2ZQY05uutXjxSaS6C3nvYfy", + "choices": [ + { + "delta": { + "content": " How" + }, + "index": 0 + } + ], + "created": 1728680302, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_e2bde53e6e" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:23.103700+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0Ik2ZQY05uutXjxSaS6C3nvYfy", + "choices": [ + { + "delta": { + "content": " can" + }, + "index": 0 + } + ], + "created": 1728680302, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_e2bde53e6e" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:23.104351+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0Ik2ZQY05uutXjxSaS6C3nvYfy", + "choices": [ + { + "delta": { + "content": " I" + }, + "index": 0 + } + ], + "created": 1728680302, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_e2bde53e6e" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:23.129299+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0Ik2ZQY05uutXjxSaS6C3nvYfy", + "choices": [ + { + "delta": { + "content": " assist" + }, + "index": 0 + } + ], + "created": 1728680302, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_e2bde53e6e" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:23.129883+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0Ik2ZQY05uutXjxSaS6C3nvYfy", + "choices": [ + { + "delta": { + "content": " you" + }, + "index": 0 + } + ], + "created": 1728680302, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_e2bde53e6e" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:23.179545+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0Ik2ZQY05uutXjxSaS6C3nvYfy", + "choices": [ + { + "delta": { + "content": " today" + }, + "index": 0 + } + ], + "created": 1728680302, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_e2bde53e6e" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:23.180217+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0Ik2ZQY05uutXjxSaS6C3nvYfy", + "choices": [ + { + "delta": { + "content": "?" + }, + "index": 0 + } + ], + "created": 1728680302, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_e2bde53e6e" + } + } + }, + { + "name": "new_token", + "time": "2024-10-11T20:58:23.180931+00:00", + "kwargs": { + "token": { + "id": "chatcmpl-AHH0Ik2ZQY05uutXjxSaS6C3nvYfy", + "choices": [ + { + "delta": {}, + "finish_reason": "stop", + "index": 0 + } + ], + "created": 1728680302, + "model": "gpt-4o-mini-2024-07-18", + "object": "chat.completion.chunk", + "system_fingerprint": "fp_e2bde53e6e" + } + } + } + ] + } + ] +} \ No newline at end of file diff --git a/python/tests/integration_tests/wrappers/test_openai.py b/python/tests/integration_tests/wrappers/test_openai.py index 32dcd85c2..2a344c966 100644 --- a/python/tests/integration_tests/wrappers/test_openai.py +++ b/python/tests/integration_tests/wrappers/test_openai.py @@ -1,5 +1,9 @@ # mypy: disable-error-code="attr-defined, union-attr, arg-type, call-overload" +import json +import os import time +from pathlib import Path +from typing import Any from unittest import mock import pytest @@ -180,3 +184,188 @@ async def test_completions_async_api(mock_session: mock.MagicMock, stream: bool) assert mock_session.return_value.request.call_count >= 1 for call in mock_session.return_value.request.call_args_list[1:]: assert call[0][0].upper() == "POST" + + +class Collect: + """ + Collects the runs for inspection. + """ + + def __init__(self): + self.run = None + + def __call__(self, run): + self.run = run + + +def _collect_requests(mock_session: mock.MagicMock, filename: str): + mock_requests = mock_session.return_value.request.call_args_list + collected_requests = {} + for _ in range(10): + time.sleep(0.1) + for call in mock_requests: + if json_bytes := call.kwargs.get("data"): + json_str = json_bytes.decode("utf-8") + collected_requests.update(json.loads(json_str)) + all_events = [ + *collected_requests.get("post", []), + *collected_requests.get("patch", []), + ] + # if end_time has been set, we can stop collecting as the background + # thread has finished processing the run + if any(event.get("end_time") for event in all_events): + break + mock_session.return_value.request.call_args_list.clear() + + if os.environ.get("WRITE_TOKEN_COUNTING_TEST_DATA") == "1": + dir_path = Path(__file__).resolve().parent.parent / "test_data" + file_path = dir_path / f"{filename}.json" + with open(file_path, "w") as f: + json.dump(collected_requests, f, indent=2) + + +test_cases = [ + { + "description": "stream", + "params": { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "howdy"}], + "stream": True, + "stream_options": {"include_usage": True}, + }, + "expect_usage_metadata": True, + }, + { + "description": "stream no usage", + "params": { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "howdy"}], + "stream": True, + }, + "expect_usage_metadata": False, + }, + { + "description": "", + "params": { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "howdy"}], + }, + "expect_usage_metadata": True, + }, + { + "description": "reasoning", + "params": { + "model": "o1-mini", + "messages": [ + { + "role": "user", + "content": ( + "Write a bash script that takes a matrix represented " + "as a string with format '[1,2],[3,4],[5,6]' and prints the " + "transpose in the same format." + ), + } + ], + }, + "expect_usage_metadata": True, + "check_reasoning_tokens": True, + }, +] + + +@pytest.mark.parametrize("test_case", test_cases) +@mock.patch("langsmith.client.requests.Session") +def test_wrap_openai_chat_tokens(mock_session: mock.MagicMock, test_case): + import openai + from openai.types.chat import ChatCompletion, ChatCompletionChunk + + oai_client = openai.Client() + ls_client = langsmith.Client(session=mock_session()) + wrapped_oai_client = wrap_openai(oai_client, tracing_extra={"client": ls_client}) + + collect = Collect() + run_id_to_usage_metadata = {} + with langsmith.tracing_context(enabled=True): + params: dict[str, Any] = test_case["params"].copy() + params["langsmith_extra"] = {"on_end": collect} + res = wrapped_oai_client.chat.completions.create(**params) + + if params.get("stream"): + for chunk in res: + assert isinstance(chunk, ChatCompletionChunk) + if test_case.get("expect_usage_metadata") and hasattr(chunk, "usage"): + oai_usage = chunk.usage + else: + assert isinstance(res, ChatCompletion) + oai_usage = res.usage + + if test_case["expect_usage_metadata"]: + usage_metadata = collect.run.outputs["usage_metadata"] + assert usage_metadata["input_tokens"] == oai_usage.prompt_tokens + assert usage_metadata["output_tokens"] == oai_usage.completion_tokens + assert usage_metadata["total_tokens"] == oai_usage.total_tokens + if test_case.get("check_reasoning_tokens"): + assert ( + usage_metadata["output_token_details"]["reasoning"] + == oai_usage.completion_tokens_details.reasoning_tokens + ) + else: + assert collect.run.outputs.get("usage_metadata") is None + assert collect.run.outputs.get("usage") is None + + run_id_to_usage_metadata[collect.run.id] = collect.run.outputs.get( + "usage_metadata" + ) + + filename = f"langsmith_py_wrap_openai_{test_case['description'].replace(' ', '_')}" + _collect_requests(mock_session, filename) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_case", test_cases) +@mock.patch("langsmith.client.requests.Session") +async def test_wrap_openai_chat_async_tokens(mock_session: mock.MagicMock, test_case): + import openai + from openai.types.chat import ChatCompletion, ChatCompletionChunk + + oai_client = openai.AsyncClient() + ls_client = langsmith.Client(session=mock_session()) + wrapped_oai_client = wrap_openai(oai_client, tracing_extra={"client": ls_client}) + + collect = Collect() + run_id_to_usage_metadata = {} + with langsmith.tracing_context(enabled=True): + params: dict[str, Any] = test_case["params"].copy() + params["langsmith_extra"] = {"on_end": collect} + res = await wrapped_oai_client.chat.completions.create(**params) + + if params.get("stream"): + oai_usage = None + async for chunk in res: + assert isinstance(chunk, ChatCompletionChunk) + if test_case.get("expect_usage_metadata") and hasattr(chunk, "usage"): + oai_usage = chunk.usage + else: + assert isinstance(res, ChatCompletion) + oai_usage = res.usage + + if test_case["expect_usage_metadata"]: + usage_metadata = collect.run.outputs["usage_metadata"] + assert usage_metadata["input_tokens"] == oai_usage.prompt_tokens + assert usage_metadata["output_tokens"] == oai_usage.completion_tokens + assert usage_metadata["total_tokens"] == oai_usage.total_tokens + if test_case.get("check_reasoning_tokens"): + assert ( + usage_metadata["output_token_details"]["reasoning"] + == oai_usage.completion_tokens_details.reasoning_tokens + ) + else: + assert collect.run.outputs.get("usage_metadata") is None + assert collect.run.outputs.get("usage") is None + + run_id_to_usage_metadata[collect.run.id] = collect.run.outputs.get( + "usage_metadata" + ) + + filename = f"langsmith_py_wrap_openai_{test_case['description'].replace(' ', '_')}" + _collect_requests(mock_session, filename)