From b4fcec73d43b4775d7d28c0fa28649290aa2fd62 Mon Sep 17 00:00:00 2001 From: Kevin Guevara Date: Thu, 2 Jan 2025 22:38:45 -0500 Subject: [PATCH 1/6] feat: mistralail v1 chat --- pyproject.toml | 7 +- scope3ai/lib.py | 9 ++ scope3ai/tracers/mistrarlai_v1/chat.py | 131 ++++++++++++++++ scope3ai/tracers/mistrarlai_v1/instrument.py | 40 +++++ .../cassettes/test_mistralai_async_chat.yaml | 79 ++++++++++ .../test_mistralai_async_stream_chat.yaml | 147 ++++++++++++++++++ tests/cassettes/test_mistralai_chat.yaml | 79 ++++++++++ .../cassettes/test_mistralai_stream_chat.yaml | 147 ++++++++++++++++++ tests/test_mistralai_v1.py | 50 ++++++ uv.lock | 147 ++++++++---------- 10 files changed, 751 insertions(+), 85 deletions(-) create mode 100644 scope3ai/tracers/mistrarlai_v1/chat.py create mode 100644 scope3ai/tracers/mistrarlai_v1/instrument.py create mode 100644 tests/cassettes/test_mistralai_async_chat.yaml create mode 100644 tests/cassettes/test_mistralai_async_stream_chat.yaml create mode 100644 tests/cassettes/test_mistralai_chat.yaml create mode 100644 tests/cassettes/test_mistralai_stream_chat.yaml create mode 100644 tests/test_mistralai_v1.py diff --git a/pyproject.toml b/pyproject.toml index dfaa142..4f5b363 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,6 @@ description = "Track the environmental impact of your use of AI" readme = "README.md" requires-python = ">=3.9" dependencies = [ - "httpx>=0.28.1", "litellm>=1.53.3", "pillow>=11.0.0", "pydantic>=2.10.3", @@ -29,9 +28,6 @@ litellm = [ "litellm>=1.53.3", "rapidfuzz>=3.10.1", ] -mistralai = [ - "mistralai>=0.4.2", -] anthropic = [ "anthropic>=0.40.0", ] @@ -44,6 +40,9 @@ huggingface-hub = [ "minijinja>=2.5.0", "tiktoken>=0.8.0", ] +mistralai = [ + "mistralai==1.2.5" +] google-generativeai = [ "google-generativeai>=0.8.3", ] diff --git a/scope3ai/lib.py b/scope3ai/lib.py index 9e57c1a..db99d13 100644 --- a/scope3ai/lib.py +++ b/scope3ai/lib.py @@ -59,12 +59,21 @@ def init_litellm_instrumentor() -> None: instrumentor.instrument() +def init_mistral_v1_instrumentor() -> None: + if importlib.util.find_spec("mistralai") is not None: + from scope3ai.tracers.mistrarlai_v1.instrument import MistralAIInstrumentor + + instrumentor = MistralAIInstrumentor() + instrumentor.instrument() + + _INSTRUMENTS = { "anthropic": init_anthropic_instrumentor, "cohere": init_cohere_instrumentor, "openai": init_openai_instrumentor, "huggingface_hub": init_huggingface_hub_instrumentor, "litellm": init_litellm_instrumentor, + "mistralai": init_mistral_v1_instrumentor, } diff --git a/scope3ai/tracers/mistrarlai_v1/chat.py b/scope3ai/tracers/mistrarlai_v1/chat.py new file mode 100644 index 0000000..587f923 --- /dev/null +++ b/scope3ai/tracers/mistrarlai_v1/chat.py @@ -0,0 +1,131 @@ +import time +from collections.abc import AsyncGenerator, Iterable +from typing import Any, Callable, Optional + +from mistralai import Mistral +from mistralai.models import ChatCompletionResponse as _ChatCompletionResponse +from mistralai.models import CompletionChunk as _CompletionChunk +from mistralai.models import CompletionEvent + +from scope3ai import Scope3AI +from scope3ai.api.types import Scope3AIContext +from scope3ai.api.typesgen import ImpactRow, Model + +PROVIDER = "mistralai" + + +class ChatCompletionResponse(_ChatCompletionResponse): + scope3ai: Optional[Scope3AIContext] = None + + +class CompletionChunk(_CompletionChunk): + scope3ai: Optional[Scope3AIContext] = None + + +def mistralai_v1_chat_wrapper( + wrapped: Callable, + instance: Mistral, + args: Any, + kwargs: Any, # noqa: ARG001 +) -> ChatCompletionResponse: + timer_start = time.perf_counter() + response = wrapped(*args, **kwargs) + request_latency = time.perf_counter() - timer_start + scope3_row = ImpactRow( + model=Model(id=response.model), + input_tokens=response.usage.prompt_tokens, + output_tokens=response.usage.completion_tokens, + request_duration_ms=request_latency * 1000, + managed_service_id=PROVIDER, + ) + scope3ai_ctx = Scope3AI.get_instance().submit_impact(scope3_row) + chat = ChatCompletionResponse(**response.model_dump()) + chat.scope3ai = scope3ai_ctx + return chat + + +def mistralai_v1_chat_wrapper_stream( + wrapped: Callable, + instance: Mistral, + args: Any, + kwargs: Any, # noqa: ARG001 +) -> Iterable[CompletionEvent]: + timer_start = time.perf_counter() + stream = wrapped(*args, **kwargs) + token_count = 0 + for i, chunk in enumerate(stream): + if i > 0 and chunk.data.choices[0].finish_reason is None: + token_count += 1 + model_name = chunk.data.model + if chunk.data: + request_latency = time.perf_counter() - timer_start + scope3_row = ImpactRow( + model=Model(id=model_name), + input_tokens=token_count, + output_tokens=chunk.data.usage.completion_tokens + if chunk.data.usage + else None, + request_duration_ms=request_latency * 1000, + managed_service_id=PROVIDER, + ) + scope3ai_ctx = Scope3AI.get_instance().submit_impact(scope3_row) + chunk.data = CompletionChunk( + **chunk.data.model_dump(), scope3ai=scope3ai_ctx + ) + yield chunk + + +async def mistralai_v1_async_chat_wrapper( + wrapped: Callable, + instance: Mistral, # noqa: ARG001 + args: Any, + kwargs: Any, +) -> ChatCompletionResponse: + timer_start = time.perf_counter() + response = await wrapped(*args, **kwargs) + request_latency = time.perf_counter() - timer_start + scope3_row = ImpactRow( + model=Model(id=response.model), + input_tokens=response.usage.prompt_tokens, + output_tokens=response.usage.completion_tokens, + request_duration_ms=request_latency * 1000, + managed_service_id=PROVIDER, + ) + scope3ai_ctx = Scope3AI.get_instance().submit_impact(scope3_row) + chat = ChatCompletionResponse(**response.model_dump()) + chat.scope3ai = scope3ai_ctx + return chat + + +async def _generator( + stream: AsyncGenerator[CompletionEvent, None], timer_start: float +) -> AsyncGenerator[CompletionEvent, None]: + token_count = 0 + async for chunk in stream: + if chunk.data.usage is not None: + token_count = chunk.data.usage.completion_tokens + request_latency = time.perf_counter() - timer_start + model_name = chunk.data.model + scope3_row = ImpactRow( + model=Model(id=model_name), + input_tokens=token_count, + output_tokens=chunk.data.usage.completion_tokens + if chunk.data.usage + else None, + request_duration_ms=request_latency * 1000, + managed_service_id=PROVIDER, + ) + scope3ai_ctx = Scope3AI.get_instance().submit_impact(scope3_row) + chunk.data = CompletionChunk(**chunk.data.model_dump(), scope3ai=scope3ai_ctx) + yield chunk + + +async def mistralai_v1_async_chat_wrapper_stream( + wrapped: Callable, + instance: Mistral, # noqa: ARG001 + args: Any, + kwargs: Any, +) -> AsyncGenerator[CompletionEvent, None]: + timer_start = time.perf_counter() + stream = await wrapped(*args, **kwargs) + return _generator(stream, timer_start) diff --git a/scope3ai/tracers/mistrarlai_v1/instrument.py b/scope3ai/tracers/mistrarlai_v1/instrument.py new file mode 100644 index 0000000..628626a --- /dev/null +++ b/scope3ai/tracers/mistrarlai_v1/instrument.py @@ -0,0 +1,40 @@ +from wrapt import wrap_function_wrapper # type: ignore[import-untyped] + +from scope3ai.tracers.mistrarlai_v1.chat import ( + mistralai_v1_chat_wrapper, + mistralai_v1_async_chat_wrapper, + mistralai_v1_chat_wrapper_stream, + mistralai_v1_async_chat_wrapper_stream, +) + + +class MistralAIInstrumentor: + def __init__(self) -> None: + self.wrapped_methods = [ + { + "module": "mistralai.chat", + "name": "Chat.complete", + "wrapper": mistralai_v1_chat_wrapper, + }, + { + "module": "mistralai.chat", + "name": "Chat.complete_async", + "wrapper": mistralai_v1_async_chat_wrapper, + }, + { + "module": "mistralai.chat", + "name": "Chat.stream", + "wrapper": mistralai_v1_chat_wrapper_stream, + }, + { + "module": "mistralai.chat", + "name": "Chat.stream_async", + "wrapper": mistralai_v1_async_chat_wrapper_stream, + }, + ] + + def instrument(self) -> None: + for wrapper in self.wrapped_methods: + wrap_function_wrapper( + wrapper["module"], wrapper["name"], wrapper["wrapper"] + ) diff --git a/tests/cassettes/test_mistralai_async_chat.yaml b/tests/cassettes/test_mistralai_async_chat.yaml new file mode 100644 index 0000000..6a07f54 --- /dev/null +++ b/tests/cassettes/test_mistralai_async_chat.yaml @@ -0,0 +1,79 @@ +interactions: +- request: + body: '{"messages":[{"content":"Hello World!","role":"user"}],"model":"mistral-tiny","safe_prompt":false,"stream":false,"top_p":1.0}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + authorization: + - DUMMY + connection: + - keep-alive + content-length: + - '125' + content-type: + - application/json + host: + - api.mistral.ai + user-agent: + - mistral-client-python/1.2.5 + method: POST + uri: https://api.mistral.ai/v1/chat/completions + response: + body: + string: !!binary | + H4sIAAAAAAAAAzzQQW7CMBCF4au4b9ONiwIEknjTLZyhqtDgTBpTxxPFg1qEcvcqqmD762n0ae4I + LRwa6qqq2NTtvi7L3fbc1Pumq6nzBXclNTUs5Hxhr3DwPenKyzBG1iAJFn5iUm7h1tV2V1fromgs + Bmk5wmEIWSeKbxrSbdn2EjxnuI87Qmr5F66wGDhn+mK4OyaJDAfKOWSlpLBQkXjyFGOGS9cYLbwk + 5bRoDhyjGO154hdz1NdsUvBsVMzArOYm15U5yI/xlMzR/F9dqlFp6faO2aILKeT+NDFlSXDIKiPm + T4vrwzROMox6UvnmlOH2C0kpPsOmXEiPjzzzup7nPwAAAP//AwBG0Z1HYQEAAA== + headers: + CF-Cache-Status: + - DYNAMIC + CF-RAY: + - 8fbf86f0889ada8f-MIA + Connection: + - keep-alive + Content-Encoding: + - gzip + Content-Type: + - application/json + Date: + - Fri, 03 Jan 2025 02:23:29 GMT + Server: + - cloudflare + Transfer-Encoding: + - chunked + access-control-allow-origin: + - '*' + alt-svc: + - h3=":443"; ma=86400 + ratelimitbysize-limit: + - '500000' + ratelimitbysize-query-cost: + - '32003' + ratelimitbysize-remaining: + - '467997' + ratelimitbysize-reset: + - '31' + x-envoy-upstream-service-time: + - '141' + x-kong-proxy-latency: + - '1' + x-kong-request-id: + - 564b303d23cf7f55aa9c761e5fa8216e + x-kong-upstream-latency: + - '142' + x-ratelimitbysize-limit-minute: + - '500000' + x-ratelimitbysize-limit-month: + - '1000000000' + x-ratelimitbysize-remaining-minute: + - '467997' + x-ratelimitbysize-remaining-month: + - '999903991' + status: + code: 200 + message: OK +version: 1 diff --git a/tests/cassettes/test_mistralai_async_stream_chat.yaml b/tests/cassettes/test_mistralai_async_stream_chat.yaml new file mode 100644 index 0000000..e4b3e20 --- /dev/null +++ b/tests/cassettes/test_mistralai_async_stream_chat.yaml @@ -0,0 +1,147 @@ +interactions: +- request: + body: '{"messages":[{"content":"Hello World!","role":"user"}],"model":"mistral-tiny","safe_prompt":false,"stream":true,"top_p":1.0}' + headers: + accept: + - text/event-stream + accept-encoding: + - gzip, deflate + authorization: + - DUMMY + connection: + - keep-alive + content-length: + - '124' + content-type: + - application/json + host: + - api.mistral.ai + user-agent: + - mistral-client-python/1.2.5 + method: POST + uri: https://api.mistral.ai/v1/chat/completions#stream + response: + body: + string: 'data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + there"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":"!"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + It"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":"''"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":"s"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + nice"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + to"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + meet"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + you"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":"."},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + How"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + can"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + I"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + help"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + you"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + today"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":"?"},"finish_reason":null}]} + + + data: {"id":"0084bd861cc84921abde79c5afc3151f","object":"chat.completion.chunk","created":1735870849,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":""},"finish_reason":"stop"}],"usage":{"prompt_tokens":6,"total_tokens":24,"completion_tokens":18}} + + + data: [DONE] + + + ' + headers: + CF-Cache-Status: + - DYNAMIC + CF-RAY: + - 8fbf83070d268de2-MIA + Connection: + - keep-alive + Content-Type: + - text/event-stream; charset=utf-8 + Date: + - Fri, 03 Jan 2025 02:20:49 GMT + Server: + - cloudflare + Transfer-Encoding: + - chunked + access-control-allow-origin: + - '*' + alt-svc: + - h3=":443"; ma=86400 + ratelimitbysize-limit: + - '500000' + ratelimitbysize-query-cost: + - '32003' + ratelimitbysize-remaining: + - '435994' + ratelimitbysize-reset: + - '11' + x-envoy-upstream-service-time: + - '49' + x-kong-proxy-latency: + - '0' + x-kong-request-id: + - b1c3b413c8044414c8b02342846f0a72 + x-kong-upstream-latency: + - '50' + x-ratelimitbysize-limit-minute: + - '500000' + x-ratelimitbysize-limit-month: + - '1000000000' + x-ratelimitbysize-remaining-minute: + - '435994' + x-ratelimitbysize-remaining-month: + - '999935994' + status: + code: 200 + message: OK +version: 1 diff --git a/tests/cassettes/test_mistralai_chat.yaml b/tests/cassettes/test_mistralai_chat.yaml new file mode 100644 index 0000000..3e3bb34 --- /dev/null +++ b/tests/cassettes/test_mistralai_chat.yaml @@ -0,0 +1,79 @@ +interactions: +- request: + body: '{"messages":[{"content":"Hello World!","role":"user"}],"model":"mistral-tiny","safe_prompt":false,"stream":false,"top_p":1.0}' + headers: + accept: + - application/json + accept-encoding: + - gzip, deflate + authorization: + - DUMMY + connection: + - keep-alive + content-length: + - '125' + content-type: + - application/json + host: + - api.mistral.ai + user-agent: + - mistral-client-python/1.2.5 + method: POST + uri: https://api.mistral.ai/v1/chat/completions + response: + body: + string: !!binary | + H4sIAAAAAAAAAzzQQWrDMBCF4auob9ONGhxbaRxtuk3OUEpQpEmtVtYYa0wbjO9eTEm2P4/hY2bE + AItAjWnr0NTmcDGNa9pAlTOHaldfqTU1QYMvX+QFFr5zsvHcD4kkcoaGH8kJBdjtvtm1+6o1e42e + AyVY9LHI6NKLxHxbtx1HTwX2fUbMgX5hK42eSnGfBDtj5ESwcKXEIi4LNIQ5nb1LqcDmKSUNz1ko + r5ojpcRKOhrpSZ3kuagcPSlh1ROJuvG0UUf+Ud5ldVL/V9eqhIO7vWHRuMYcS3ceyRXOsCjCA5YP + jeluGkbuBzkLf1MusK8rSVx6hNqspPtHHnnbLssfAAAA//8DADq2FIdhAQAA + headers: + CF-Cache-Status: + - DYNAMIC + CF-RAY: + - 8fbf82fae98e4c2c-MIA + Connection: + - keep-alive + Content-Encoding: + - gzip + Content-Type: + - application/json + Date: + - Fri, 03 Jan 2025 02:20:47 GMT + Server: + - cloudflare + Transfer-Encoding: + - chunked + access-control-allow-origin: + - '*' + alt-svc: + - h3=":443"; ma=86400 + ratelimitbysize-limit: + - '500000' + ratelimitbysize-query-cost: + - '32003' + ratelimitbysize-remaining: + - '467997' + ratelimitbysize-reset: + - '13' + x-envoy-upstream-service-time: + - '136' + x-kong-proxy-latency: + - '0' + x-kong-request-id: + - 42cab944bfba9d617b23e97d80914930 + x-kong-upstream-latency: + - '137' + x-ratelimitbysize-limit-minute: + - '500000' + x-ratelimitbysize-limit-month: + - '1000000000' + x-ratelimitbysize-remaining-minute: + - '467997' + x-ratelimitbysize-remaining-month: + - '999967997' + status: + code: 200 + message: OK +version: 1 diff --git a/tests/cassettes/test_mistralai_stream_chat.yaml b/tests/cassettes/test_mistralai_stream_chat.yaml new file mode 100644 index 0000000..77ce19a --- /dev/null +++ b/tests/cassettes/test_mistralai_stream_chat.yaml @@ -0,0 +1,147 @@ +interactions: +- request: + body: '{"messages":[{"content":"Hello World!","role":"user"}],"model":"mistral-tiny","safe_prompt":false,"stream":true,"top_p":1.0}' + headers: + accept: + - text/event-stream + accept-encoding: + - gzip, deflate + authorization: + - DUMMY + connection: + - keep-alive + content-length: + - '124' + content-type: + - application/json + host: + - api.mistral.ai + user-agent: + - mistral-client-python/1.2.5 + method: POST + uri: https://api.mistral.ai/v1/chat/completions#stream + response: + body: + string: 'data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":"Hello"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + there"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":"!"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + It"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":"''"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":"s"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + nice"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + to"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + meet"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + you"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":"."},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + How"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + can"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + I"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + assist"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + you"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":" + today"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":"?"},"finish_reason":null}]} + + + data: {"id":"0687886dfb4045f59ae9e817a1d4dfe1","object":"chat.completion.chunk","created":1735871117,"model":"mistral-tiny","choices":[{"index":0,"delta":{"content":""},"finish_reason":"stop"}],"usage":{"prompt_tokens":6,"total_tokens":24,"completion_tokens":18}} + + + data: [DONE] + + + ' + headers: + CF-Cache-Status: + - DYNAMIC + CF-RAY: + - 8fbf89915a8edaed-MIA + Connection: + - keep-alive + Content-Type: + - text/event-stream; charset=utf-8 + Date: + - Fri, 03 Jan 2025 02:25:17 GMT + Server: + - cloudflare + Transfer-Encoding: + - chunked + access-control-allow-origin: + - '*' + alt-svc: + - h3=":443"; ma=86400 + ratelimitbysize-limit: + - '500000' + ratelimitbysize-query-cost: + - '32003' + ratelimitbysize-remaining: + - '467997' + ratelimitbysize-reset: + - '43' + x-envoy-upstream-service-time: + - '76' + x-kong-proxy-latency: + - '4' + x-kong-request-id: + - ee0f4c550debd85688d6c047842c687d + x-kong-upstream-latency: + - '77' + x-ratelimitbysize-limit-minute: + - '500000' + x-ratelimitbysize-limit-month: + - '1000000000' + x-ratelimitbysize-remaining-minute: + - '467997' + x-ratelimitbysize-remaining-month: + - '999903991' + status: + code: 200 + message: OK +version: 1 diff --git a/tests/test_mistralai_v1.py b/tests/test_mistralai_v1.py new file mode 100644 index 0000000..a8d1457 --- /dev/null +++ b/tests/test_mistralai_v1.py @@ -0,0 +1,50 @@ +import pytest +from mistralai import Mistral + + +@pytest.mark.vcr +def test_mistralai_chat(tracer_init): + client = Mistral() + response = client.chat.complete( + messages=[{"role": "user", "content": "Hello World!"}], model="mistral-tiny" + ) + assert len(response.choices) > 0 + assert getattr(response, "scope3ai") is not None + assert response.scope3ai.request.input_tokens == 6 + assert response.scope3ai.request.output_tokens == 18 + assert response.scope3ai.impact is None + + +@pytest.mark.vcr +@pytest.mark.asyncio +async def test_mistralai_async_chat(tracer_init): + client = Mistral() + response = await client.chat.complete_async( + messages=[{"role": "user", "content": "Hello World!"}], model="mistral-tiny" + ) + assert len(response.choices) > 0 + assert getattr(response, "scope3ai") is not None + assert response.scope3ai.request.input_tokens == 6 + assert response.scope3ai.request.output_tokens == 18 + assert response.scope3ai.impact is None + + +@pytest.mark.vcr +def test_mistralai_stream_chat(tracer_init): + client = Mistral() + stream = client.chat.stream( + messages=[{"role": "user", "content": "Hello World!"}], model="mistral-tiny" + ) + for chunk in stream: + assert getattr(chunk.data, "scope3ai") is not None + + +@pytest.mark.vcr +@pytest.mark.asyncio +async def test_mistralai_async_stream_chat(tracer_init): + client = Mistral() + stream = await client.chat.stream_async( + messages=[{"role": "user", "content": "Hello World!"}], model="mistral-tiny" + ) + async for chunk in stream: + assert getattr(chunk.data, "scope3ai") is not None diff --git a/uv.lock b/uv.lock index fb207c2..60e8098 100644 --- a/uv.lock +++ b/uv.lock @@ -343,7 +343,7 @@ name = "click" version = "8.1.7" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "colorama", marker = "platform_system == 'Windows'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/96/d3/f04c7bfcf5c1862a2a5b845c6b2b360488cf47af55dfa79c98f6a6bf98b5/click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de", size = 336121 } wheels = [ @@ -434,6 +434,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d7/ee/bf0adb559ad3c786f12bcbc9296b3f5675f529199bef03e2df281fa1fadb/email_validator-2.2.0-py3-none-any.whl", hash = "sha256:561977c2d73ce3611850a06fa56b414621e0c8faa9d66f2611407d87465da631", size = 33521 }, ] +[[package]] +name = "eval-type-backport" +version = "0.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/30/ea/8b0ac4469d4c347c6a385ff09dc3c048c2d021696664e26c7ee6791631b5/eval_type_backport-0.2.2.tar.gz", hash = "sha256:f0576b4cf01ebb5bd358d02314d31846af5e07678387486e2c798af0e7d849c1", size = 9079 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/31/55cd413eaccd39125368be33c46de24a1f639f2e12349b0361b4678f3915/eval_type_backport-0.2.2-py3-none-any.whl", hash = "sha256:cb6ad7c393517f476f96d456d0412ea80f0a8cf96f6892834cd9340149111b0a", size = 5830 }, +] + [[package]] name = "exceptiongroup" version = "1.2.2" @@ -799,17 +808,18 @@ wheels = [ [[package]] name = "httpx" -version = "0.28.1" +version = "0.27.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, { name = "certifi" }, { name = "httpcore" }, { name = "idna" }, + { name = "sniffio" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406 } +sdist = { url = "https://files.pythonhosted.org/packages/78/82/08f8c936781f67d9e6b9eeb8a0c8b4e406136ea4c3d1f89a5db71d42e0e6/httpx-0.27.2.tar.gz", hash = "sha256:f7c2be1d2f3c3c3160d441802406b206c2b76f5947b11115e6df10c6c65e66c2", size = 144189 } wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517 }, + { url = "https://files.pythonhosted.org/packages/56/95/9377bcb415797e44274b51d46e3249eba641711cf3348050f76ee7b15ffc/httpx-0.27.2-py3-none-any.whl", hash = "sha256:7bb2708e112d8fdd7829cd4243970f0c223274051cb35ee80c03301ee29a3df0", size = 76395 }, ] [[package]] @@ -970,6 +980,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/32/b7/a3cde72c644fd1caf9da07fb38cf2c130f43484d8f91011940b7c4f42c8f/jiter-0.8.2-cp39-cp39-win_amd64.whl", hash = "sha256:1c0dfbd1be3cbefc7510102370d86e35d1d53e5a93d48519688b1bf0f761160a", size = 207527 }, ] +[[package]] +name = "jsonpath-python" +version = "1.0.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b5/49/e582e50b0c54c1b47e714241c4a4767bf28758bf90212248aea8e1ce8516/jsonpath-python-1.0.6.tar.gz", hash = "sha256:dd5be4a72d8a2995c3f583cf82bf3cd1a9544cfdabf2d22595b67aff07349666", size = 18121 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/16/8a/d63959f4eff03893a00e6e63592e3a9f15b9266ed8e0275ab77f8c7dbc94/jsonpath_python-1.0.6-py3-none-any.whl", hash = "sha256:1e3b78df579f5efc23565293612decee04214609208a2335884b3ee3f786b575", size = 7552 }, +] + [[package]] name = "jsonschema" version = "4.23.0" @@ -1104,16 +1123,19 @@ wheels = [ [[package]] name = "mistralai" -version = "0.4.2" +version = "1.2.5" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "eval-type-backport" }, { name = "httpx" }, - { name = "orjson" }, + { name = "jsonpath-python" }, { name = "pydantic" }, + { name = "python-dateutil" }, + { name = "typing-inspect" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/fa/20/4204f461588310b3a7ffbbbb7fa573493dc1c8185d376ee72516c04575bf/mistralai-0.4.2.tar.gz", hash = "sha256:5eb656710517168ae053f9847b0bb7f617eda07f1f93f946ad6c91a4d407fd93", size = 14234 } +sdist = { url = "https://files.pythonhosted.org/packages/55/34/95efe73fd3cd0d5f3f0198b2bfc570dfe485aa5045100aa97fa176dcb653/mistralai-1.2.5.tar.gz", hash = "sha256:05d4130f79704e3b19c0b6320944a348547879fce6894feeb72d9e9d0ee65151", size = 132348 } wheels = [ - { url = "https://files.pythonhosted.org/packages/4f/fe/79dad76b8d94b62d9e2aab8446183190e1dc384c617d06c3c93307850e11/mistralai-0.4.2-py3-none-any.whl", hash = "sha256:63c98eea139585f0a3b2c4c6c09c453738bac3958055e6f2362d3866e96b0168", size = 20334 }, + { url = "https://files.pythonhosted.org/packages/85/08/279a3afe0b319c283ae6d1ee8d42c606855093579e93e51cce2f6ced91a7/mistralai-1.2.5-py3-none-any.whl", hash = "sha256:5f0ef2680ead0329569111face1bf2ff7c67c454d43aa0e21324a8faf6c3ab22", size = 260045 }, ] [[package]] @@ -1274,75 +1296,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6b/f0/c25aab1845afc3583a808d0a9b844bb3842245d87b9791e588e6aba8ce3a/openai-1.57.1-py3-none-any.whl", hash = "sha256:3865686c927e93492d1145938d4a24b634951531c4b2769d43ca5dbd4b25d8fd", size = 389846 }, ] -[[package]] -name = "orjson" -version = "3.10.12" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e0/04/bb9f72987e7f62fb591d6c880c0caaa16238e4e530cbc3bdc84a7372d75f/orjson-3.10.12.tar.gz", hash = "sha256:0a78bbda3aea0f9f079057ee1ee8a1ecf790d4f1af88dd67493c6b8ee52506ff", size = 5438647 } -wheels = [ - { url = "https://files.pythonhosted.org/packages/72/d2/78652b67f86d093dca984ce3fa5bf819ee1462627da83e7d0b784a9a7c45/orjson-3.10.12-cp310-cp310-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:ece01a7ec71d9940cc654c482907a6b65df27251255097629d0dea781f255c6d", size = 248688 }, - { url = "https://files.pythonhosted.org/packages/70/cb/f8b6a52f3bc724edf8a62d8d1d8ee17cf19d6ae1cac89f077f0e7c30f396/orjson-3.10.12-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c34ec9aebc04f11f4b978dd6caf697a2df2dd9b47d35aa4cc606cabcb9df69d7", size = 136952 }, - { url = "https://files.pythonhosted.org/packages/a6/43/c55700df9814545bc8c35d87395ec4b9ee473a3c1f5ed72f8d3ad0298ee9/orjson-3.10.12-cp310-cp310-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:fd6ec8658da3480939c79b9e9e27e0db31dffcd4ba69c334e98c9976ac29140e", size = 149089 }, - { url = "https://files.pythonhosted.org/packages/07/da/e7e7d73bd971710b736fbd8330b8830c5fa4fc0ac003b31af61f03b26dfc/orjson-3.10.12-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f17e6baf4cf01534c9de8a16c0c611f3d94925d1701bf5f4aff17003677d8ced", size = 140479 }, - { url = "https://files.pythonhosted.org/packages/08/49/c9dfddba56ff24eecfacf2f01a76cae4d249ac2995b1359bf63a74b1b318/orjson-3.10.12-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:6402ebb74a14ef96f94a868569f5dccf70d791de49feb73180eb3c6fda2ade56", size = 156564 }, - { url = "https://files.pythonhosted.org/packages/96/df/174d2eff227dc23b4540a0c2efa6ec8fe406c442c4b7f0f556242f026d1f/orjson-3.10.12-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0000758ae7c7853e0a4a6063f534c61656ebff644391e1f81698c1b2d2fc8cd2", size = 131282 }, - { url = "https://files.pythonhosted.org/packages/6a/96/8628c53a52e2a0a1ee861d809092df72aabbd312c71de9ad6d49e2c039ab/orjson-3.10.12-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:888442dcee99fd1e5bd37a4abb94930915ca6af4db50e23e746cdf4d1e63db13", size = 139764 }, - { url = "https://files.pythonhosted.org/packages/38/17/08becb49e59e7bb7b29dc1dad19bc0c48635e627ee27e60eb5b64efcf7b1/orjson-3.10.12-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:c1f7a3ce79246aa0e92f5458d86c54f257fb5dfdc14a192651ba7ec2c00f8a05", size = 131913 }, - { url = "https://files.pythonhosted.org/packages/2a/05/f32acc2500e3fafee9445eb8b2a6ff19c4641035e6059c6c8d7bdb3abc9e/orjson-3.10.12-cp310-cp310-musllinux_1_2_armv7l.whl", hash = "sha256:802a3935f45605c66fb4a586488a38af63cb37aaad1c1d94c982c40dcc452e85", size = 415782 }, - { url = "https://files.pythonhosted.org/packages/06/03/6cc740d998d8bb60e75d4b7e228d18964475239ac842cc1865d49d092545/orjson-3.10.12-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:1da1ef0113a2be19bb6c557fb0ec2d79c92ebd2fed4cfb1b26bab93f021fb885", size = 142383 }, - { url = "https://files.pythonhosted.org/packages/f8/30/39cac82547fe021615376245c558b216d3ae8c99bd6b2274f312e49f1c94/orjson-3.10.12-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:7a3273e99f367f137d5b3fecb5e9f45bcdbfac2a8b2f32fbc72129bbd48789c2", size = 130661 }, - { url = "https://files.pythonhosted.org/packages/95/29/c6837f4fc1eaa742eaf5abcd767ab6805493f44fe1f72b37c1743706c1d8/orjson-3.10.12-cp310-none-win32.whl", hash = "sha256:475661bf249fd7907d9b0a2a2421b4e684355a77ceef85b8352439a9163418c3", size = 143625 }, - { url = "https://files.pythonhosted.org/packages/f6/62/c6b955f2144421108fa441b5471e1d5f8654a7df9840b261106e04d5d15c/orjson-3.10.12-cp310-none-win_amd64.whl", hash = "sha256:87251dc1fb2b9e5ab91ce65d8f4caf21910d99ba8fb24b49fd0c118b2362d509", size = 135075 }, - { url = "https://files.pythonhosted.org/packages/d3/48/7c3cd094488f5a3bc58488555244609a8c4d105bc02f2b77e509debf0450/orjson-3.10.12-cp311-cp311-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:a734c62efa42e7df94926d70fe7d37621c783dea9f707a98cdea796964d4cf74", size = 248687 }, - { url = "https://files.pythonhosted.org/packages/ff/90/e55f0e25c7fdd1f82551fe787f85df6f378170caca863c04c810cd8f2730/orjson-3.10.12-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:750f8b27259d3409eda8350c2919a58b0cfcd2054ddc1bd317a643afc646ef23", size = 136953 }, - { url = "https://files.pythonhosted.org/packages/2a/b3/109c020cf7fee747d400de53b43b183ca9d3ebda3906ad0b858eb5479718/orjson-3.10.12-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:bb52c22bfffe2857e7aa13b4622afd0dd9d16ea7cc65fd2bf318d3223b1b6252", size = 149090 }, - { url = "https://files.pythonhosted.org/packages/96/d4/35c0275dc1350707d182a1b5da16d1184b9439848060af541285407f18f9/orjson-3.10.12-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:440d9a337ac8c199ff8251e100c62e9488924c92852362cd27af0e67308c16ef", size = 140480 }, - { url = "https://files.pythonhosted.org/packages/3b/79/f863ff460c291ad2d882cc3b580cc444bd4ec60c9df55f6901e6c9a3f519/orjson-3.10.12-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a9e15c06491c69997dfa067369baab3bf094ecb74be9912bdc4339972323f252", size = 156564 }, - { url = "https://files.pythonhosted.org/packages/98/7e/8d5835449ddd873424ee7b1c4ba73a0369c1055750990d824081652874d6/orjson-3.10.12-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:362d204ad4b0b8724cf370d0cd917bb2dc913c394030da748a3bb632445ce7c4", size = 131279 }, - { url = "https://files.pythonhosted.org/packages/46/f5/d34595b6d7f4f984c6fef289269a7f98abcdc2445ebdf90e9273487dda6b/orjson-3.10.12-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:2b57cbb4031153db37b41622eac67329c7810e5f480fda4cfd30542186f006ae", size = 139764 }, - { url = "https://files.pythonhosted.org/packages/b3/5b/ee6e9ddeab54a7b7806768151c2090a2d36025bc346a944f51cf172ef7f7/orjson-3.10.12-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:165c89b53ef03ce0d7c59ca5c82fa65fe13ddf52eeb22e859e58c237d4e33b9b", size = 131915 }, - { url = "https://files.pythonhosted.org/packages/c4/45/febee5951aef6db5cd8cdb260548101d7ece0ca9d4ddadadf1766306b7a4/orjson-3.10.12-cp311-cp311-musllinux_1_2_armv7l.whl", hash = "sha256:5dee91b8dfd54557c1a1596eb90bcd47dbcd26b0baaed919e6861f076583e9da", size = 415783 }, - { url = "https://files.pythonhosted.org/packages/27/a5/5a8569e49f3a6c093bee954a3de95062a231196f59e59df13a48e2420081/orjson-3.10.12-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:77a4e1cfb72de6f905bdff061172adfb3caf7a4578ebf481d8f0530879476c07", size = 142387 }, - { url = "https://files.pythonhosted.org/packages/6e/05/02550fb38c5bf758f3994f55401233a2ef304e175f473f2ac6dbf464cc8b/orjson-3.10.12-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:038d42c7bc0606443459b8fe2d1f121db474c49067d8d14c6a075bbea8bf14dd", size = 130664 }, - { url = "https://files.pythonhosted.org/packages/8c/f4/ba31019d0646ce51f7ac75af6dabf98fd89dbf8ad87a9086da34710738e7/orjson-3.10.12-cp311-none-win32.whl", hash = "sha256:03b553c02ab39bed249bedd4abe37b2118324d1674e639b33fab3d1dafdf4d79", size = 143623 }, - { url = "https://files.pythonhosted.org/packages/83/fe/babf08842b989acf4c46103fefbd7301f026423fab47e6f3ba07b54d7837/orjson-3.10.12-cp311-none-win_amd64.whl", hash = "sha256:8b8713b9e46a45b2af6b96f559bfb13b1e02006f4242c156cbadef27800a55a8", size = 135074 }, - { url = "https://files.pythonhosted.org/packages/a1/2f/989adcafad49afb535da56b95d8f87d82e748548b2a86003ac129314079c/orjson-3.10.12-cp312-cp312-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:53206d72eb656ca5ac7d3a7141e83c5bbd3ac30d5eccfe019409177a57634b0d", size = 248678 }, - { url = "https://files.pythonhosted.org/packages/69/b9/8c075e21a50c387649db262b618ebb7e4d40f4197b949c146fc225dd23da/orjson-3.10.12-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ac8010afc2150d417ebda810e8df08dd3f544e0dd2acab5370cfa6bcc0662f8f", size = 136763 }, - { url = "https://files.pythonhosted.org/packages/87/d3/78edf10b4ab14c19f6d918cf46a145818f4aca2b5a1773c894c5490d3a4c/orjson-3.10.12-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:ed459b46012ae950dd2e17150e838ab08215421487371fa79d0eced8d1461d70", size = 149137 }, - { url = "https://files.pythonhosted.org/packages/16/81/5db8852bdf990a0ddc997fa8f16b80895b8cc77c0fe3701569ed2b4b9e78/orjson-3.10.12-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8dcb9673f108a93c1b52bfc51b0af422c2d08d4fc710ce9c839faad25020bb69", size = 140567 }, - { url = "https://files.pythonhosted.org/packages/fa/a6/9ce1e3e3db918512efadad489630c25841eb148513d21dab96f6b4157fa1/orjson-3.10.12-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:22a51ae77680c5c4652ebc63a83d5255ac7d65582891d9424b566fb3b5375ee9", size = 156620 }, - { url = "https://files.pythonhosted.org/packages/47/d4/05133d6bea24e292d2f7628b1e19986554f7d97b6412b3e51d812e38db2d/orjson-3.10.12-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:910fdf2ac0637b9a77d1aad65f803bac414f0b06f720073438a7bd8906298192", size = 131555 }, - { url = "https://files.pythonhosted.org/packages/b9/7a/b3fbffda8743135c7811e95dc2ab7cdbc5f04999b83c2957d046f1b3fac9/orjson-3.10.12-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:24ce85f7100160936bc2116c09d1a8492639418633119a2224114f67f63a4559", size = 139743 }, - { url = "https://files.pythonhosted.org/packages/b5/13/95bbcc9a6584aa083da5ce5004ce3d59ea362a542a0b0938d884fd8790b6/orjson-3.10.12-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8a76ba5fc8dd9c913640292df27bff80a685bed3a3c990d59aa6ce24c352f8fc", size = 131733 }, - { url = "https://files.pythonhosted.org/packages/e8/29/dddbb2ea6e7af426fcc3da65a370618a88141de75c6603313d70768d1df1/orjson-3.10.12-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:ff70ef093895fd53f4055ca75f93f047e088d1430888ca1229393a7c0521100f", size = 415788 }, - { url = "https://files.pythonhosted.org/packages/53/df/4aea59324ac539975919b4705ee086aced38e351a6eb3eea0f5071dd5661/orjson-3.10.12-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:f4244b7018b5753ecd10a6d324ec1f347da130c953a9c88432c7fbc8875d13be", size = 142347 }, - { url = "https://files.pythonhosted.org/packages/55/55/a52d83d7c49f8ff44e0daab10554490447d6c658771569e1c662aa7057fe/orjson-3.10.12-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:16135ccca03445f37921fa4b585cff9a58aa8d81ebcb27622e69bfadd220b32c", size = 130829 }, - { url = "https://files.pythonhosted.org/packages/a1/8b/b1beb1624dd4adf7d72e2d9b73c4b529e7851c0c754f17858ea13e368b33/orjson-3.10.12-cp312-none-win32.whl", hash = "sha256:2d879c81172d583e34153d524fcba5d4adafbab8349a7b9f16ae511c2cee8708", size = 143659 }, - { url = "https://files.pythonhosted.org/packages/13/91/634c9cd0bfc6a857fc8fab9bf1a1bd9f7f3345e0d6ca5c3d4569ceb6dcfa/orjson-3.10.12-cp312-none-win_amd64.whl", hash = "sha256:fc23f691fa0f5c140576b8c365bc942d577d861a9ee1142e4db468e4e17094fb", size = 135221 }, - { url = "https://files.pythonhosted.org/packages/1b/bb/3f560735f46fa6f875a9d7c4c2171a58cfb19f56a633d5ad5037a924f35f/orjson-3.10.12-cp313-cp313-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:47962841b2a8aa9a258b377f5188db31ba49af47d4003a32f55d6f8b19006543", size = 248662 }, - { url = "https://files.pythonhosted.org/packages/a3/df/54817902350636cc9270db20486442ab0e4db33b38555300a1159b439d16/orjson-3.10.12-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6334730e2532e77b6054e87ca84f3072bee308a45a452ea0bffbbbc40a67e296", size = 126055 }, - { url = "https://files.pythonhosted.org/packages/2e/77/55835914894e00332601a74540840f7665e81f20b3e2b9a97614af8565ed/orjson-3.10.12-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:accfe93f42713c899fdac2747e8d0d5c659592df2792888c6c5f829472e4f85e", size = 131507 }, - { url = "https://files.pythonhosted.org/packages/33/9e/b91288361898e3158062a876b5013c519a5d13e692ac7686e3486c4133ab/orjson-3.10.12-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a7974c490c014c48810d1dede6c754c3cc46598da758c25ca3b4001ac45b703f", size = 131686 }, - { url = "https://files.pythonhosted.org/packages/b2/15/08ce117d60a4d2d3fd24e6b21db463139a658e9f52d22c9c30af279b4187/orjson-3.10.12-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:3f250ce7727b0b2682f834a3facff88e310f52f07a5dcfd852d99637d386e79e", size = 415710 }, - { url = "https://files.pythonhosted.org/packages/71/af/c09da5ed58f9c002cf83adff7a4cdf3e6cee742aa9723395f8dcdb397233/orjson-3.10.12-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:f31422ff9486ae484f10ffc51b5ab2a60359e92d0716fcce1b3593d7bb8a9af6", size = 142305 }, - { url = "https://files.pythonhosted.org/packages/17/d1/8612038d44f33fae231e9ba480d273bac2b0383ce9e77cb06bede1224ae3/orjson-3.10.12-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:5f29c5d282bb2d577c2a6bbde88d8fdcc4919c593f806aac50133f01b733846e", size = 130815 }, - { url = "https://files.pythonhosted.org/packages/67/2c/d5f87834be3591555cfaf9aecdf28f480a6f0b4afeaac53bad534bf9518f/orjson-3.10.12-cp313-none-win32.whl", hash = "sha256:f45653775f38f63dc0e6cd4f14323984c3149c05d6007b58cb154dd080ddc0dc", size = 143664 }, - { url = "https://files.pythonhosted.org/packages/6a/05/7d768fa3ca23c9b3e1e09117abeded1501119f1d8de0ab722938c91ab25d/orjson-3.10.12-cp313-none-win_amd64.whl", hash = "sha256:229994d0c376d5bdc91d92b3c9e6be2f1fbabd4cc1b59daae1443a46ee5e9825", size = 134944 }, - { url = "https://files.pythonhosted.org/packages/01/50/5a52cfe65fc70e7b843cbe143b850313095d4e45f99aeb278b4b3691f3cf/orjson-3.10.12-cp39-cp39-macosx_10_15_x86_64.macosx_11_0_arm64.macosx_10_15_universal2.whl", hash = "sha256:f29de3ef71a42a5822765def1febfb36e0859d33abf5c2ad240acad5c6a1b78d", size = 249130 }, - { url = "https://files.pythonhosted.org/packages/af/42/adb00fc60890e57ee6022257d70d9a20dfd28bfd955401e2d02a60192226/orjson-3.10.12-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:de365a42acc65d74953f05e4772c974dad6c51cfc13c3240899f534d611be967", size = 136784 }, - { url = "https://files.pythonhosted.org/packages/f5/f6/5ef130a2e28a0c633c82da67224212fa84b17f93d9d768c255f389d66641/orjson-3.10.12-cp39-cp39-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:91a5a0158648a67ff0004cb0df5df7dcc55bfc9ca154d9c01597a23ad54c8d0c", size = 148936 }, - { url = "https://files.pythonhosted.org/packages/54/80/f40805bb094d3470bcfca8d30ae9f4606bb90c809190fee18c17a4653956/orjson-3.10.12-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c47ce6b8d90fe9646a25b6fb52284a14ff215c9595914af63a5933a49972ce36", size = 140291 }, - { url = "https://files.pythonhosted.org/packages/b4/0a/c91265a075af28fcb118a9a690dccc0f31c52dd7f486f5cc73aa6f9a69b4/orjson-3.10.12-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0eee4c2c5bfb5c1b47a5db80d2ac7aaa7e938956ae88089f098aff2c0f35d5d8", size = 156335 }, - { url = "https://files.pythonhosted.org/packages/6f/1f/c8fc64b25e2e4a5fa0fb514a0325719e99461c088a3b659954a57b8c1f3c/orjson-3.10.12-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:35d3081bbe8b86587eb5c98a73b97f13d8f9fea685cf91a579beddacc0d10566", size = 131118 }, - { url = "https://files.pythonhosted.org/packages/b3/f7/f5dba457bbc6aaf8ebd0131a17eb672835d2ea31f763d1462c191902bc10/orjson-3.10.12-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:73c23a6e90383884068bc2dba83d5222c9fcc3b99a0ed2411d38150734236755", size = 139573 }, - { url = "https://files.pythonhosted.org/packages/85/03/87091a1c831076e8d6ed0be19d30bd837d3c27e29f43b6b2dc23efee3d29/orjson-3.10.12-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:5472be7dc3269b4b52acba1433dac239215366f89dc1d8d0e64029abac4e714e", size = 131760 }, - { url = "https://files.pythonhosted.org/packages/53/70/956ca7999fe43b5e2ca51248f732850acaaf00f8fa5b8f7924bc504157bb/orjson-3.10.12-cp39-cp39-musllinux_1_2_armv7l.whl", hash = "sha256:7319cda750fca96ae5973efb31b17d97a5c5225ae0bc79bf5bf84df9e1ec2ab6", size = 415605 }, - { url = "https://files.pythonhosted.org/packages/ca/7f/4f49c1a9da03e383198c0fa418cc4262d01d70931742dfc504c2a495ed7b/orjson-3.10.12-cp39-cp39-musllinux_1_2_i686.whl", hash = "sha256:74d5ca5a255bf20b8def6a2b96b1e18ad37b4a122d59b154c458ee9494377f80", size = 142208 }, - { url = "https://files.pythonhosted.org/packages/81/d9/0216c950e295b1c493510ef546b80328dc06e2ae4a82bc693386e4cdebeb/orjson-3.10.12-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:ff31d22ecc5fb85ef62c7d4afe8301d10c558d00dd24274d4bbe464380d3cd69", size = 130495 }, - { url = "https://files.pythonhosted.org/packages/1a/d0/70d97b583d628a307d4bc1f96f7ab6cf137e32849c06ea03d4337ad19e14/orjson-3.10.12-cp39-none-win32.whl", hash = "sha256:c22c3ea6fba91d84fcb4cda30e64aff548fcf0c44c876e681f47d61d24b12e6b", size = 143467 }, - { url = "https://files.pythonhosted.org/packages/ec/a3/725e4cac70d2a88f5fc4f9eed451e12d594063823feb931c30ac1394d26f/orjson-3.10.12-cp39-none-win_amd64.whl", hash = "sha256:be604f60d45ace6b0b33dd990a66b4526f1a7a186ac411c942674625456ca548", size = 134927 }, -] - [[package]] name = "packaging" version = "24.2" @@ -1788,6 +1741,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9d/d3/ff520d11e6ee400602711d1ece8168dcfc5b6d8146fb7db4244a6ad6a9c3/pytest_vcr-1.0.2-py2.py3-none-any.whl", hash = "sha256:2f316e0539399bea0296e8b8401145c62b6f85e9066af7e57b6151481b0d6d9c", size = 4137 }, ] +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892 }, +] + [[package]] name = "python-dotenv" version = "1.0.1" @@ -2186,7 +2151,6 @@ name = "scope3ai" version = "0.2.1" source = { virtual = "." } dependencies = [ - { name = "httpx" }, { name = "litellm" }, { name = "pillow" }, { name = "pydantic" }, @@ -2237,12 +2201,11 @@ requires-dist = [ { name = "anthropic", marker = "extra == 'anthropic'", specifier = ">=0.40.0" }, { name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.13.3" }, { name = "google-generativeai", marker = "extra == 'google-generativeai'", specifier = ">=0.8.3" }, - { name = "httpx", specifier = ">=0.28.1" }, { name = "huggingface-hub", marker = "extra == 'huggingface-hub'", specifier = ">=0.26.5" }, { name = "litellm", specifier = ">=1.53.3" }, { name = "litellm", marker = "extra == 'litellm'", specifier = ">=1.53.3" }, { name = "minijinja", marker = "extra == 'huggingface-hub'", specifier = ">=2.5.0" }, - { name = "mistralai", marker = "extra == 'mistralai'", specifier = ">=0.4.2" }, + { name = "mistralai", marker = "extra == 'mistralai'", specifier = "==1.2.5" }, { name = "openai", marker = "extra == 'openai'", specifier = ">=1.57.1" }, { name = "pillow", specifier = ">=11.0.0" }, { name = "pydantic", specifier = ">=2.10.3" }, @@ -2260,6 +2223,15 @@ tests = [ { name = "pytest-vcr", specifier = ">=1.0.2" }, ] +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050 }, +] + [[package]] name = "sniffio" version = "1.3.1" @@ -2389,7 +2361,7 @@ name = "tqdm" version = "4.67.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "colorama", marker = "platform_system == 'Windows'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737 } wheels = [ @@ -2454,6 +2426,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/26/9f/ad63fc0248c5379346306f8668cda6e2e2e9c95e01216d2b8ffd9ff037d0/typing_extensions-4.12.2-py3-none-any.whl", hash = "sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d", size = 37438 }, ] +[[package]] +name = "typing-inspect" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mypy-extensions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/dc/74/1789779d91f1961fa9438e9a8710cdae6bd138c80d7303996933d117264a/typing_inspect-0.9.0.tar.gz", hash = "sha256:b23fc42ff6f6ef6954e4852c1fb512cdd18dbea03134f91f856a95ccc9461f78", size = 13825 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/65/f3/107a22063bf27bdccf2024833d3445f4eea42b2e598abfbd46f6a63b6cb0/typing_inspect-0.9.0-py3-none-any.whl", hash = "sha256:9ee6fc59062311ef8547596ab6b955e1b8aa46242d854bfc78f4f6b0eff35f9f", size = 8827 }, +] + [[package]] name = "uritemplate" version = "4.1.1" From 7482f84c102cc7bf5993ebf5efa64a6bc8c41817 Mon Sep 17 00:00:00 2001 From: Kevin Guevara Date: Thu, 2 Jan 2025 22:45:56 -0500 Subject: [PATCH 2/6] fix: remove some console logs --- scope3ai/tracers/mistrarlai_v1/chat.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/scope3ai/tracers/mistrarlai_v1/chat.py b/scope3ai/tracers/mistrarlai_v1/chat.py index 587f923..cd21a43 100644 --- a/scope3ai/tracers/mistrarlai_v1/chat.py +++ b/scope3ai/tracers/mistrarlai_v1/chat.py @@ -26,7 +26,7 @@ def mistralai_v1_chat_wrapper( wrapped: Callable, instance: Mistral, args: Any, - kwargs: Any, # noqa: ARG001 + kwargs: Any, ) -> ChatCompletionResponse: timer_start = time.perf_counter() response = wrapped(*args, **kwargs) @@ -48,7 +48,7 @@ def mistralai_v1_chat_wrapper_stream( wrapped: Callable, instance: Mistral, args: Any, - kwargs: Any, # noqa: ARG001 + kwargs: Any, ) -> Iterable[CompletionEvent]: timer_start = time.perf_counter() stream = wrapped(*args, **kwargs) @@ -77,7 +77,7 @@ def mistralai_v1_chat_wrapper_stream( async def mistralai_v1_async_chat_wrapper( wrapped: Callable, - instance: Mistral, # noqa: ARG001 + instance: Mistral, args: Any, kwargs: Any, ) -> ChatCompletionResponse: @@ -122,7 +122,7 @@ async def _generator( async def mistralai_v1_async_chat_wrapper_stream( wrapped: Callable, - instance: Mistral, # noqa: ARG001 + instance: Mistral, args: Any, kwargs: Any, ) -> AsyncGenerator[CompletionEvent, None]: From 844ad53078987b045e065251827c3c8b18484014 Mon Sep 17 00:00:00 2001 From: Kevin Guevara Date: Thu, 2 Jan 2025 22:48:03 -0500 Subject: [PATCH 3/6] fix: add mistral to readme --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index e2abf21..b133929 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,7 @@ uv add scope3ai | OpenAI | ✅ | | | | | | Huggingface | ✅ | ✅ | ✅ | ✅ | ✅ | | LiteLLM | ✅ | | | | | +| MistralAi | ✅ | | | | | Roadmap: - Google From 8422c638564d61fb98e91911c7c81bc5603b4589 Mon Sep 17 00:00:00 2001 From: Kevin Guevara Date: Thu, 2 Jan 2025 22:58:49 -0500 Subject: [PATCH 4/6] feat: adding async tests --- tests/test_huggingface_tracer.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_huggingface_tracer.py b/tests/test_huggingface_tracer.py index 8ddc6a9..9bbf78e 100644 --- a/tests/test_huggingface_tracer.py +++ b/tests/test_huggingface_tracer.py @@ -63,6 +63,14 @@ def test_huggingface_hub_image_generation(tracer_init): assert response.scope3ai.request.request_duration_ms == pytest.approx(18850, 0.1) +@pytest.mark.vcr +@pytest.mark.asyncio +async def test_huggingface_hub_image_generation_async(tracer_init): + client = AsyncInferenceClient() + image = await client.text_to_image("An astronaut riding a horse on the moon.") + assert image + + @pytest.mark.vcr def test_huggingface_hub_translation(tracer_init): client = InferenceClient() From fb3bb7a78c5cfefddfc2010fb10b47bea05a4a55 Mon Sep 17 00:00:00 2001 From: Kevin Guevara Date: Fri, 3 Jan 2025 11:07:22 -0500 Subject: [PATCH 5/6] feat: async capture response --- scope3ai/tracers/huggingface/instrument.py | 16 +++++++-- scope3ai/tracers/huggingface/text_to_image.py | 34 +++++-------------- scope3ai/tracers/huggingface/utils.py | 34 +++++++++++++++++-- 3 files changed, 53 insertions(+), 31 deletions(-) diff --git a/scope3ai/tracers/huggingface/instrument.py b/scope3ai/tracers/huggingface/instrument.py index c9099f9..a0dc0a6 100644 --- a/scope3ai/tracers/huggingface/instrument.py +++ b/scope3ai/tracers/huggingface/instrument.py @@ -7,7 +7,10 @@ from scope3ai.tracers.huggingface.speech_to_text import ( huggingface_automatic_recognition_output_wrapper, ) -from scope3ai.tracers.huggingface.text_to_image import huggingface_text_to_image_wrapper +from scope3ai.tracers.huggingface.text_to_image import ( + huggingface_text_to_image_wrapper, + huggingface_text_to_image_wrapper_async, +) from scope3ai.tracers.huggingface.text_to_speech import ( huggingface_text_to_speech_wrapper, ) @@ -16,8 +19,9 @@ ) from .utils import ( hf_raise_for_status_enabled, - get_client_session_async, hf_raise_for_status_wrapper, + hf_async_raise_for_status_enabled, + get_client_session_async_wrapper, ) @@ -63,7 +67,13 @@ def __init__(self) -> None: { "module": "huggingface_hub.inference._generated._async_client", "name": "AsyncInferenceClient._get_client_session", - "wrapper": get_client_session_async, + "wrapper": get_client_session_async_wrapper, + "enable": hf_async_raise_for_status_enabled, + }, + { + "module": "huggingface_hub.inference._generated._async_client", + "name": "AsyncInferenceClient.text_to_image", + "wrapper": huggingface_text_to_image_wrapper_async, }, ] diff --git a/scope3ai/tracers/huggingface/text_to_image.py b/scope3ai/tracers/huggingface/text_to_image.py index 93d199f..3a05314 100644 --- a/scope3ai/tracers/huggingface/text_to_image.py +++ b/scope3ai/tracers/huggingface/text_to_image.py @@ -8,7 +8,10 @@ from scope3ai.api.types import Scope3AIContext, Model, ImpactRow from scope3ai.api.typesgen import Task from scope3ai.lib import Scope3AI -from scope3ai.tracers.huggingface.utils import hf_raise_for_status_capture +from scope3ai.tracers.huggingface.utils import ( + hf_raise_for_status_capture, + hf_async_raise_for_status_capture, +) PROVIDER = "huggingface_hub" @@ -51,32 +54,11 @@ def huggingface_text_to_image_wrapper_non_stream( async def huggingface_text_to_image_wrapper_async_non_stream( wrapped: Callable, instance: AsyncInferenceClient, args: Any, kwargs: Any ) -> TextToImageOutput: - response = await wrapped(*args, **kwargs) - with hf_raise_for_status_capture() as capture_response: - response = wrapped(*args, **kwargs) + with hf_async_raise_for_status_capture() as capture_response: + response = await wrapped(*args, **kwargs) http_response = capture_response.get() - model = kwargs.get("model") or instance.get_recommended_model("text-to-speech") - encoder = tiktoken.get_encoding("cl100k_base") - if len(args) > 0: - prompt = args[0] - else: - prompt = kwargs["prompt"] - compute_time = http_response.headers.get("x-compute-time") - input_tokens = len(encoder.encode(prompt)) - width, height = response.size - scope3_row = ImpactRow( - model=Model(id=model), - input_tokens=input_tokens, - task=Task.text_to_image, - output_images=["{width}x{height}".format(width=width, height=height)], - request_duration_ms=float(compute_time) * 1000, - managed_service_id=PROVIDER, - ) - - scope3_ctx = Scope3AI.get_instance().submit_impact(scope3_row) - result = TextToImageOutput(response) - result.scope3ai = scope3_ctx - return result + print(http_response, "RESPONSE IS CAPTURED SUCCESSFULLY") + return response def huggingface_text_to_image_wrapper( diff --git a/scope3ai/tracers/huggingface/utils.py b/scope3ai/tracers/huggingface/utils.py index 434a360..35b8b5d 100644 --- a/scope3ai/tracers/huggingface/utils.py +++ b/scope3ai/tracers/huggingface/utils.py @@ -6,10 +6,21 @@ HFRS_VALUE = contextvars.ContextVar(f"{HFRS_BASEKEY}__value", default=None) +HFRS_ASYNC_BASEKEY = "scope3ai__huggingface__hf_async_raise_for_status" +HFRS_ASYNC_ENABLED = contextvars.ContextVar( + f"{HFRS_ASYNC_BASEKEY}__enabled", default=None +) +HFRS_ASYNC_VALUE = contextvars.ContextVar(f"{HFRS_ASYNC_BASEKEY}__value", default=None) + + def hf_raise_for_status_enabled(): return HFRS_ENABLED.get() is True +def hf_async_raise_for_status_enabled(): + return HFRS_ASYNC_ENABLED.get() is True + + def hf_raise_for_status_wrapper(wrapped, instance, args, kwargs): try: result = wrapped(*args, **kwargs) @@ -30,10 +41,29 @@ def hf_raise_for_status_capture(): HFRS_ENABLED.set(False) -def get_client_session_async(wrapped, instance, args, kwargs): +@contextlib.contextmanager +def hf_async_raise_for_status_capture(): + try: + HFRS_ASYNC_VALUE.set(None) + HFRS_ASYNC_ENABLED.set(True) + yield HFRS_ASYNC_VALUE + finally: + HFRS_ASYNC_ENABLED.set(False) + + +def async_post_wrapper(session_post): + async def wrapped_post(*args, **kwargs): + result = await session_post(*args, **kwargs) + HFRS_ASYNC_VALUE.set(result) + return result + + return wrapped_post + + +def get_client_session_async_wrapper(wrapped, instance, args, kwargs): try: result = wrapped(*args, **kwargs) - print("session", result.__dir__()) + result.post = async_post_wrapper(result.post) return result except Exception as e: raise e From 0248fef90c40885a78dfec6729edf7cde3029ee1 Mon Sep 17 00:00:00 2001 From: Kevin Guevara Date: Fri, 3 Jan 2025 11:35:34 -0500 Subject: [PATCH 6/6] fix: text-to-image async implementation --- scope3ai/tracers/huggingface/text_to_image.py | 24 +++++++++++++++++-- tests/test_huggingface_tracer.py | 9 +++++-- 2 files changed, 29 insertions(+), 4 deletions(-) diff --git a/scope3ai/tracers/huggingface/text_to_image.py b/scope3ai/tracers/huggingface/text_to_image.py index 3a05314..df98ab1 100644 --- a/scope3ai/tracers/huggingface/text_to_image.py +++ b/scope3ai/tracers/huggingface/text_to_image.py @@ -57,8 +57,28 @@ async def huggingface_text_to_image_wrapper_async_non_stream( with hf_async_raise_for_status_capture() as capture_response: response = await wrapped(*args, **kwargs) http_response = capture_response.get() - print(http_response, "RESPONSE IS CAPTURED SUCCESSFULLY") - return response + model = kwargs.get("model") or instance.get_recommended_model("text-to-speech") + encoder = tiktoken.get_encoding("cl100k_base") + if len(args) > 0: + prompt = args[0] + else: + prompt = kwargs["prompt"] + compute_time = http_response.headers.get("x-compute-time") + input_tokens = len(encoder.encode(prompt)) + width, height = response.size + scope3_row = ImpactRow( + model=Model(id=model), + input_tokens=input_tokens, + task=Task.text_to_image, + output_images=["{width}x{height}".format(width=width, height=height)], + request_duration_ms=float(compute_time) * 1000, + managed_service_id=PROVIDER, + ) + + scope3_ctx = Scope3AI.get_instance().submit_impact(scope3_row) + result = TextToImageOutput(response) + result.scope3ai = scope3_ctx + return result def huggingface_text_to_image_wrapper( diff --git a/tests/test_huggingface_tracer.py b/tests/test_huggingface_tracer.py index b5fb7dc..dbdd1b8 100644 --- a/tests/test_huggingface_tracer.py +++ b/tests/test_huggingface_tracer.py @@ -70,8 +70,13 @@ def test_huggingface_hub_image_generation(tracer_init): @pytest.mark.asyncio async def test_huggingface_hub_image_generation_async(tracer_init): client = AsyncInferenceClient() - image = await client.text_to_image("An astronaut riding a horse on the moon.") - assert image + response = await client.text_to_image("An astronaut riding a horse on the moon.") + assert response.image + assert getattr(response, "scope3ai") is not None + assert response.scope3ai.request.input_tokens == 9 + assert len(response.scope3ai.request.output_images) == 1 + assert response.scope3ai.impact is None + assert response.scope3ai.request.request_duration_ms == pytest.approx(18850, 0.1) @pytest.mark.vcr