From 8bda6d499a701db8c3aeefc9d41a358750a6ea49 Mon Sep 17 00:00:00 2001 From: Brendan Date: Thu, 6 Jun 2024 11:53:44 -0700 Subject: [PATCH 1/3] initial otel llm specs implemented --- mirascope/anthropic/types.py | 10 ++ mirascope/base/types.py | 30 +++++ mirascope/cohere/types.py | 10 ++ mirascope/gemini/types.py | 25 ++++ mirascope/groq/types.py | 10 ++ mirascope/logfire/logfire.py | 2 + mirascope/mistral/types.py | 10 ++ mirascope/openai/types.py | 20 ++++ mirascope/otel/otel.py | 217 +++++++++++++++++++---------------- 9 files changed, 237 insertions(+), 97 deletions(-) diff --git a/mirascope/anthropic/types.py b/mirascope/anthropic/types.py index 4c3e2658..6757de80 100644 --- a/mirascope/anthropic/types.py +++ b/mirascope/anthropic/types.py @@ -147,6 +147,16 @@ def content(self) -> str: block = self.response.content[0] return block.text if block.type == "text" else "" + @property + def id(self) -> str: + """Returns the id of the response.""" + return self.response.id + + @property + def finish_reasons(self) -> Optional[list[str]]: + """Returns the finish reason of the response.""" + return [self.response.stop_reason] + @property def usage(self) -> Usage: """Returns the usage of the message.""" diff --git a/mirascope/base/types.py b/mirascope/base/types.py index 668e0251..5f82c016 100644 --- a/mirascope/base/types.py +++ b/mirascope/base/types.py @@ -163,6 +163,21 @@ def content(self) -> str: """ ... # pragma: no cover + @property + @abstractmethod + def finish_reasons(self) -> list[str]: + """Should return the finish reasons of the response. + + If there is no finish reason, this method must return None. + """ + ... # pragma: no cover + + @property + @abstractmethod + def id(self) -> Optional[str]: + """Should return the id of the response.""" + ... # pragma: no cover + @property @abstractmethod def usage(self) -> Any: @@ -218,3 +233,18 @@ def content(self) -> str: the empty string. """ ... # pragma: no cover + + @property + @abstractmethod + def finish_reasons(self) -> list[str]: + """Should return the finish reasons of the response. + + If there is no finish reason, this method must return None. + """ + ... # pragma: no cover + + @property + @abstractmethod + def id(self) -> Optional[str]: + """Should return the id of the response.""" + ... # pragma: no cover diff --git a/mirascope/cohere/types.py b/mirascope/cohere/types.py index 31f4bafe..5956dd51 100644 --- a/mirascope/cohere/types.py +++ b/mirascope/cohere/types.py @@ -121,6 +121,16 @@ def content(self) -> str: """Returns the content of the chat completion for the 0th choice.""" return self.response.text + @property + def id(self) -> Optional[str]: + """Returns the id of the response.""" + return self.response.generation_id + + @property + def finish_reasons(self) -> Optional[list[str]]: + """Returns the finish reasons of the response.""" + return [self.response.finish_reason] + @property def search_queries(self) -> Optional[list[ChatSearchQuery]]: """Returns the search queries for the 0th choice message.""" diff --git a/mirascope/gemini/types.py b/mirascope/gemini/types.py index 7952682d..4dc9cdff 100644 --- a/mirascope/gemini/types.py +++ b/mirascope/gemini/types.py @@ -120,6 +120,31 @@ def content(self) -> str: """Returns the contained string content for the 0th choice.""" return self.response.candidates[0].content.parts[0].text + @property + def id(self) -> Optional[str]: + """Returns the id of the response. + + google.generativeai does not return an id + """ + return None + + @property + def finish_reasons(self) -> list[str]: + """Returns the finish reasons of the response.""" + finish_reasons = [ + "FINISH_REASON_UNSPECIFIED", + "STOP", + "MAX_TOKENS", + "SAFETY", + "RECITATION", + "OTHER", + ] + + return [ + finish_reasons[candidate.finish_reason] + for candidate in self.response.candidates + ] + @property def usage(self) -> None: """Returns the usage of the chat completion. diff --git a/mirascope/groq/types.py b/mirascope/groq/types.py index dc23a62e..f1bea50c 100644 --- a/mirascope/groq/types.py +++ b/mirascope/groq/types.py @@ -188,6 +188,16 @@ def tool(self) -> Optional[GroqTool]: return tools[0] return None + @property + def id(self) -> str: + """Returns the id of the response.""" + return self.response.id + + @property + def finish_reasons(self) -> list[str]: + """Returns the finish reasons of the response.""" + return [choice.finish_reason for choice in self.choices] + @property def usage(self) -> Optional[Usage]: """Returns the usage of the chat completion.""" diff --git a/mirascope/logfire/logfire.py b/mirascope/logfire/logfire.py index 1de4cb92..53054bda 100644 --- a/mirascope/logfire/logfire.py +++ b/mirascope/logfire/logfire.py @@ -4,6 +4,7 @@ contextmanager, ) from string import Formatter +from textwrap import dedent from typing import Any, Callable, Optional, Union, overload import logfire @@ -243,6 +244,7 @@ def handle_before_call(self: BaseModel, fn: Callable, **kwargs): for _, var, _, _ in Formatter().parse(self.prompt_template) if var is not None } + class_vars["prompt_template"] = dedent(self.prompt_template) span_data = { "class_vars": class_vars, "template_variables": template_variables, diff --git a/mirascope/mistral/types.py b/mirascope/mistral/types.py index 0651f02c..06096eb9 100644 --- a/mirascope/mistral/types.py +++ b/mirascope/mistral/types.py @@ -91,6 +91,16 @@ def content(self) -> str: # return the first item in the list return content if isinstance(content, str) else content[0] + @property + def id(self) -> str: + """Returns the id of the response.""" + return self.response.id + + @property + def finish_reasons(self) -> list[Optional[str]]: + """Returns the finish reasons of the response.""" + return [choice.finish_reason for choice in self.choices] + @property def tool_calls(self) -> Optional[list[ToolCall]]: """Returns the tool calls for the 0th choice message.""" diff --git a/mirascope/openai/types.py b/mirascope/openai/types.py index 07673a5e..e7fb128f 100644 --- a/mirascope/openai/types.py +++ b/mirascope/openai/types.py @@ -124,6 +124,16 @@ def content(self) -> str: """Returns the content of the chat completion for the 0th choice.""" return self.message.content if self.message.content is not None else "" + @property + def id(self) -> str: + """Returns the id of the response.""" + return self.response.id + + @property + def finish_reasons(self) -> list[str]: + """Returns the finish reasons of the response.""" + return [choice.finish_reason for choice in self.response.choices] + @property def tool_calls(self) -> Optional[list[ChatCompletionMessageToolCall]]: """Returns the tool calls for the 0th choice message.""" @@ -282,6 +292,16 @@ def content(self) -> str: self.delta.content if self.delta is not None and self.delta.content else "" ) + @property + def id(self) -> str: + """Returns the id of the response.""" + return self.chunk.id + + @property + def finish_reasons(self) -> list[str]: + """Returns the finish reasons of the response.""" + return [choice.finish_reason for choice in self.chunk.choices] + @property def tool_calls(self) -> Optional[list[ChoiceDeltaToolCall]]: """Returns the partial tool calls for the 0th choice message. diff --git a/mirascope/otel/otel.py b/mirascope/otel/otel.py index 651693e7..9b77a321 100644 --- a/mirascope/otel/otel.py +++ b/mirascope/otel/otel.py @@ -14,6 +14,7 @@ SimpleSpanProcessor, ) from opentelemetry.trace import ( + SpanKind, Tracer, get_tracer, get_tracer_provider, @@ -44,9 +45,119 @@ ) -def mirascope_otel() -> Callable: +def _set_response_attributes( + span: Span, response: BaseCallResponse, messages: list[dict[str, Any]] +) -> None: + """Sets the response attributes""" + response_attributes = { + "gen_ai.response.id": response.id, + "gen_ai.response.model": response.content, + "gen_ai.response.finish_reasons": response.finish_reasons, + "gen_ai.usage.completion_tokens": response.output_tokens, + "gen_ai.usage.prompt_tokens": response.input_tokens, + } + span.set_attributes(response_attributes) + event_attributes = {"gen_ai.completion": json.dumps(messages)} + span.add_event( + "gen_ai.content.completion", + attributes=event_attributes, + ) + + +def mirascope_otel(cls) -> Callable: tracer = get_tracer("otel") + def _set_span_data( + span: Span, + suffix: str, + is_async: bool, + model_name: Optional[str], + args: tuple[Any], + ) -> None: + prompt = args[0] + if suffix == "gemini": + prompt = { + "messages": [ + {"role": message["role"], "content": message["parts"][0]} + for message in prompt + ] + } + + max_tokens = None + temperature = None + top_p = None + if hasattr(cls, "call_params"): + call_params = cls.call_params + max_tokens = getattr(call_params, "max_tokens", None) + temperature = getattr(call_params, "temperature", None) + top_p = getattr(call_params, "top_p", None) + attributes = { + "async": is_async, + "gen_ai.request.model": model_name, + "gen_ai.system": suffix, + "gen_ai.request.max_tokens": max_tokens, + "gen_ai.request.temperature": temperature, + "gen_ai.request.top_p": top_p, + } + events = { + "gen_ai.prompt": json.dumps(prompt), + } + span.set_attributes(attributes) + span.add_event("gen_ai.content.prompt", attributes=events) + + @contextmanager + def _mirascope_llm_span( + fn: Callable, + suffix: str, + is_async: bool, + model_name: Optional[str], + args: tuple[Any], + kwargs: dict[str, Any], + ): + """Wraps a pydantic class method with a OTel span.""" + tracer = get_tracer("otel") + model = kwargs.get("model", "") or model_name + with tracer.start_as_current_span( + f"{suffix}.{fn.__name__} with {model}", kind=SpanKind.CLIENT + ) as span: + _set_span_data(span, suffix, is_async, model, args) + yield span + + @contextmanager + def record_streaming( + span: Span, + content_from_stream: Callable[ + [ChunkT, type[BaseCallResponseChunk]], Union[str, None] + ], + ): + """OTel record_streaming with Mirascope providers""" + content: list[str] = [] + + def record_chunk( + chunk: ChunkT, response_chunk_type: type[BaseCallResponseChunk] + ) -> Any: + """Handles all provider chunk_types instead of only OpenAI""" + chunk_content = content_from_stream(chunk, response_chunk_type) + if chunk_content is not None: + content.append(chunk_content) + + try: + yield record_chunk + finally: + attributes = { + "response_data": { + "combined_chunk_content": "".join(content), + "chunk_count": len(content), + }, + } + span.set_attributes(attributes) + + def _extract_chunk_content( + chunk: ChunkT, response_chunk_type: type[BaseCallResponseChunk] + ) -> str: + """Extracts the content from a chunk.""" + return response_chunk_type(chunk=chunk).content + def mirascope_otel_decorator( fn: Callable, suffix: str, @@ -83,9 +194,7 @@ def wrapper(*args, **kwargs): ] message["tool_calls"] = tool_calls message.pop("content") - span.set_attribute( - "response_data", json.dumps({"message": message}) - ) + _set_response_attributes(span, response, [message]) return result async def wrapper_async(*args, **kwargs): @@ -112,20 +221,16 @@ async def wrapper_async(*args, **kwargs): ] message["tool_calls"] = tool_calls message.pop("content") - span.set_attribute( - "response_data", json.dumps({"message": message}) - ) + _set_response_attributes(span, response, [message]) return result def wrapper_generator(*args, **kwargs): - span_data = _get_span_data(suffix, is_async, model_name, args, kwargs) model = kwargs.get("model", "") or model_name with tracer.start_as_current_span( f"{suffix}.{fn.__name__} with {model}" ) as span: - with record_streaming( - span, span_data, _extract_chunk_content - ) as record_chunk: + _set_span_data(span, suffix, is_async, model, args) + with record_streaming(span, _extract_chunk_content) as record_chunk: stream = fn(*args, **kwargs) if isinstance(stream, AbstractContextManager): with stream as s: @@ -138,14 +243,12 @@ def wrapper_generator(*args, **kwargs): yield chunk async def wrapper_generator_async(*args, **kwargs): - span_data = _get_span_data(suffix, is_async, model_name, args, kwargs) model = kwargs.get("model", "") or model_name with tracer.start_as_current_span( f"{suffix}.{fn.__name__} with {model}" ) as span: - with record_streaming( - span, span_data, _extract_chunk_content - ) as record_chunk: + _set_span_data(span, suffix, is_async, model, args) + with record_streaming(span, _extract_chunk_content) as record_chunk: stream = fn(*args, **kwargs) if inspect.iscoroutine(stream): stream = await stream @@ -172,86 +275,6 @@ async def wrapper_generator_async(*args, **kwargs): return mirascope_otel_decorator -def _extract_chunk_content( - chunk: ChunkT, response_chunk_type: type[BaseCallResponseChunk] -) -> str: - """Extracts the content from a chunk.""" - return response_chunk_type(chunk=chunk).content - - -def _get_span_data( - suffix: str, - is_async: bool, - model_name: Optional[str], - args: tuple[Any], - kwargs: dict[str, Any], -) -> dict[str, Any]: - additional_request_data = {} - if suffix == "gemini": - gemini_messages = args[0] - additional_request_data = { - "messages": [ - {"role": message["role"], "content": message["parts"][0]} - for message in gemini_messages - ], - "model": model_name, - } - return { - "async": is_async, - "request_data": json.dumps(kwargs | additional_request_data), - } - - -@contextmanager -def _mirascope_llm_span( - fn: Callable, - suffix: str, - is_async: bool, - model_name: Optional[str], - args: tuple[Any], - kwargs: dict[str, Any], -): - """Wraps a pydantic class method with a Logfire span.""" - tracer = get_tracer("otel") - model = kwargs.get("model", "") or model_name - span_data = _get_span_data(suffix, is_async, model, args, kwargs) - with tracer.start_as_current_span(f"{suffix}.{fn.__name__} with {model}") as span: - span.set_attributes(span_data) - yield span - - -@contextmanager -def record_streaming( - span: Span, - span_data: dict[str, Any], - content_from_stream: Callable[ - [ChunkT, type[BaseCallResponseChunk]], Union[str, None] - ], -): - """Logfire record_streaming with Mirascope providers""" - content: list[str] = [] - - def record_chunk( - chunk: ChunkT, response_chunk_type: type[BaseCallResponseChunk] - ) -> Any: - """Handles all provider chunk_types instead of only OpenAI""" - chunk_content = content_from_stream(chunk, response_chunk_type) - if chunk_content is not None: - content.append(chunk_content) - - try: - yield record_chunk - finally: - attributes = { - **span_data, - "response_data": { - "combined_chunk_content": "".join(content), - "chunk_count": len(content), - }, - } - span.set_attributes(attributes) - - @contextmanager def handle_before_call(self: BaseModel, fn: Callable, **kwargs): """Handles before call""" @@ -330,7 +353,7 @@ def with_otel(cls: type[BaseEmbedderT]) -> type[BaseEmbedderT]: def with_otel(cls): - """Wraps a pydantic class with a Logfire span.""" + """Wraps a pydantic class with a OTel span.""" provider = get_tracer_provider() if not isinstance(provider, TracerProvider): @@ -342,6 +365,6 @@ def with_otel(cls): ) if hasattr(cls, "configuration"): cls.configuration = cls.configuration.model_copy( - update={"llm_ops": [*cls.configuration.llm_ops, mirascope_otel()]} + update={"llm_ops": [*cls.configuration.llm_ops, mirascope_otel(cls)]} ) return cls From 1a8620e7245157635bd2b4dac35a32a60a26c587 Mon Sep 17 00:00:00 2001 From: Brendan Date: Thu, 6 Jun 2024 15:48:45 -0700 Subject: [PATCH 2/3] added new properties and added tests for new properties --- mirascope/anthropic/types.py | 50 +++++++++++++++- mirascope/base/types.py | 24 ++++++-- mirascope/cohere/types.py | 55 +++++++++++++++++- mirascope/gemini/types.py | 59 +++++++++++++++++++ mirascope/groq/types.py | 41 ++++++++++++- mirascope/mistral/types.py | 45 ++++++++++++++- mirascope/openai/types.py | 14 ++++- mirascope/otel/otel.py | 105 +++++++++++++++++++++------------- tests/anthropic/conftest.py | 1 + tests/anthropic/test_types.py | 26 +++++++++ tests/cohere/test_types.py | 13 ++++- tests/conftest.py | 3 +- tests/gemini/conftest.py | 14 ++++- tests/gemini/test_types.py | 20 ++++++- tests/groq/conftest.py | 2 + tests/groq/test_types.py | 21 +++++++ tests/mistral/conftest.py | 6 +- tests/mistral/test_types.py | 13 +++++ tests/openai/test_types.py | 6 ++ tests/otel/test_otel.py | 6 +- 20 files changed, 460 insertions(+), 64 deletions(-) diff --git a/mirascope/anthropic/types.py b/mirascope/anthropic/types.py index ff27f2b8..24f55660 100644 --- a/mirascope/anthropic/types.py +++ b/mirascope/anthropic/types.py @@ -8,6 +8,7 @@ ContentBlockDeltaEvent, ContentBlockStartEvent, Message, + MessageStartEvent, MessageStreamEvent, TextBlock, TextDelta, @@ -147,6 +148,11 @@ def content(self) -> str: block = self.response.content[0] return block.text if block.type == "text" else "" + @property + def model(self) -> str: + """Returns the name of the response model.""" + return self.response.model + @property def id(self) -> str: """Returns the id of the response.""" @@ -155,7 +161,7 @@ def id(self) -> str: @property def finish_reasons(self) -> Optional[list[str]]: """Returns the finish reason of the response.""" - return [self.response.stop_reason] + return [str(self.response.stop_reason)] @property def usage(self) -> Usage: @@ -245,3 +251,45 @@ def content(self) -> str: self.chunk.delta.text if isinstance(self.chunk.delta, TextDelta) else "" ) return "" + + @property + def model(self) -> str: + """Returns the name of the response model.""" + if isinstance(self.chunk, MessageStartEvent): + return self.chunk.message.model + return None + + @property + def id(self) -> Optional[str]: + """Returns the id of the response.""" + if isinstance(self.chunk, MessageStartEvent): + return self.chunk.message.id + return None + + @property + def finish_reasons(self) -> Optional[list[str]]: + """Returns the finish reason of the response.""" + if isinstance(self.chunk, MessageStartEvent): + return [self.chunk.message.stop_reason] + return None + + @property + def usage(self) -> Optional[Usage]: + """Returns the usage of the message.""" + if isinstance(self.chunk, MessageStartEvent): + return self.chunk.message.usage + return None + + @property + def input_tokens(self) -> Optional[int]: + """Returns the number of input tokens.""" + if self.usage: + return self.usage.input_tokens + return None + + @property + def output_tokens(self) -> Optional[int]: + """Returns the number of output tokens.""" + if self.usage: + return self.usage.output_tokens + return None diff --git a/mirascope/base/types.py b/mirascope/base/types.py index 5f82c016..df752cae 100644 --- a/mirascope/base/types.py +++ b/mirascope/base/types.py @@ -165,13 +165,19 @@ def content(self) -> str: @property @abstractmethod - def finish_reasons(self) -> list[str]: + def finish_reasons(self) -> Union[None, list[str]]: """Should return the finish reasons of the response. If there is no finish reason, this method must return None. """ ... # pragma: no cover + @property + @abstractmethod + def model(self) -> Optional[str]: + """Should return the name of the response model.""" + ... # pragma: no cover + @property @abstractmethod def id(self) -> Optional[str]: @@ -236,11 +242,8 @@ def content(self) -> str: @property @abstractmethod - def finish_reasons(self) -> list[str]: - """Should return the finish reasons of the response. - - If there is no finish reason, this method must return None. - """ + def model(self) -> Optional[str]: + """Should return the name of the response model.""" ... # pragma: no cover @property @@ -248,3 +251,12 @@ def finish_reasons(self) -> list[str]: def id(self) -> Optional[str]: """Should return the id of the response.""" ... # pragma: no cover + + @property + @abstractmethod + def finish_reasons(self) -> Union[None, list[str]]: + """Should return the finish reasons of the response. + + If there is no finish reason, this method must return None. + """ + ... # pragma: no cover diff --git a/mirascope/cohere/types.py b/mirascope/cohere/types.py index de96a090..b4f7e6d2 100644 --- a/mirascope/cohere/types.py +++ b/mirascope/cohere/types.py @@ -7,6 +7,7 @@ StreamedChatResponse_SearchQueriesGeneration, StreamedChatResponse_SearchResults, StreamedChatResponse_StreamEnd, + StreamedChatResponse_StreamStart, StreamedChatResponse_ToolCallsGeneration, ) from cohere.types import ( @@ -121,6 +122,14 @@ def content(self) -> str: """Returns the content of the chat completion for the 0th choice.""" return self.response.text + @property + def model(self) -> str: + """Returns the name of the response model. + + Cohere does not return model, so we return None + """ + return None + @property def id(self) -> Optional[str]: """Returns the id of the response.""" @@ -129,7 +138,7 @@ def id(self) -> Optional[str]: @property def finish_reasons(self) -> Optional[list[str]]: """Returns the finish reasons of the response.""" - return [self.response.finish_reason] + return [str(self.response.finish_reason)] @property def search_queries(self) -> Optional[list[ChatSearchQuery]]: @@ -308,6 +317,28 @@ def citations(self) -> Optional[list[ChatCitation]]: return self.chunk.citations return None + @property + def model(self) -> str: + """Returns the name of the response model. + + Cohere does not return model, so we return None + """ + return None + + @property + def id(self) -> Optional[str]: + """Returns the id of the response.""" + if isinstance(self.chunk, StreamedChatResponse_StreamStart): + return self.chunk.generation_id + return None + + @property + def finish_reasons(self) -> Optional[list[str]]: + """Returns the finish reasons of the response.""" + if isinstance(self.chunk, StreamedChatResponse_StreamEnd): + return [str(self.chunk.finish_reason)] + return None + @property def response(self) -> Optional[NonStreamedChatResponse]: """Returns the full response for the stream-end event type else None.""" @@ -322,6 +353,28 @@ def tool_calls(self) -> Optional[list[ToolCall]]: return self.chunk.tool_calls return None + @property + def usage(self) -> Optional[ApiMetaBilledUnits]: + """Returns the usage of the response.""" + if isinstance(self.chunk, StreamedChatResponse_StreamEnd): + if self.chunk.response.meta: + return self.chunk.response.meta.billed_units + return None + + @property + def input_tokens(self) -> Optional[float]: + """Returns the number of input tokens.""" + if self.usage: + return self.usage.input_tokens + return None + + @property + def output_tokens(self) -> Optional[float]: + """Returns the number of output tokens.""" + if self.usage: + return self.usage.output_tokens + return None + class CohereEmbeddingResponse(BaseEmbeddingResponse[SkipValidation[EmbedResponse]]): """A convenience wrapper around the Cohere `EmbedResponse` response.""" diff --git a/mirascope/gemini/types.py b/mirascope/gemini/types.py index 4dc9cdff..899687c8 100644 --- a/mirascope/gemini/types.py +++ b/mirascope/gemini/types.py @@ -145,6 +145,14 @@ def finish_reasons(self) -> list[str]: for candidate in self.response.candidates ] + @property + def model(self) -> None: + """Returns the model name. + + google.generativeai does not return model, so we return None + """ + return None + @property def usage(self) -> None: """Returns the usage of the chat completion. @@ -211,3 +219,54 @@ class Math(GeminiCall): def content(self) -> str: """Returns the chunk content for the 0th choice.""" return self.chunk.candidates[0].content.parts[0].text + + @property + def id(self) -> Optional[str]: + """Returns the id of the response. + + google.generativeai does not return an id + """ + return None + + @property + def finish_reasons(self) -> list[str]: + """Returns the finish reasons of the response.""" + finish_reasons = [ + "FINISH_REASON_UNSPECIFIED", + "STOP", + "MAX_TOKENS", + "SAFETY", + "RECITATION", + "OTHER", + ] + + return [ + finish_reasons[candidate.finish_reason] + for candidate in self.chunk.candidates + ] + + @property + def model(self) -> None: + """Returns the model name. + + google.generativeai does not return model, so we return None + """ + return None + + @property + def usage(self) -> None: + """Returns the usage of the chat completion. + + google.generativeai does not have Usage, so we return None + """ + return None + + @property + def input_tokens(self) -> None: + """Returns the number of input tokens.""" + return None + + @property + def output_tokens(self) -> None: + """Returns the number of output tokens.""" + return None diff --git a/mirascope/groq/types.py b/mirascope/groq/types.py index bf6efd45..f4cdc2ca 100644 --- a/mirascope/groq/types.py +++ b/mirascope/groq/types.py @@ -193,6 +193,11 @@ def tool(self) -> Optional[GroqTool]: return tools[0] return None + @property + def model(self) -> str: + """Returns the name of the response model.""" + return self.response.model + @property def id(self) -> str: """Returns the id of the response.""" @@ -201,7 +206,7 @@ def id(self) -> str: @property def finish_reasons(self) -> list[str]: """Returns the finish reasons of the response.""" - return [choice.finish_reason for choice in self.choices] + return [str(choice.finish_reason) for choice in self.choices] @property def usage(self) -> Optional[CompletionUsage]: @@ -286,6 +291,21 @@ def content(self) -> str: """Returns the content for the 0th choice delta.""" return self.delta.content if self.delta.content is not None else "" + @property + def model(self) -> str: + """Returns the name of the response model.""" + return self.chunk.model + + @property + def id(self) -> str: + """Returns the id of the response.""" + return self.chunk.id + + @property + def finish_reasons(self) -> list[str]: + """Returns the finish reasons of the response.""" + return [str(choice.finish_reason) for choice in self.chunk.choices] + @property def tool_calls(self) -> Optional[list[ChoiceDeltaToolCall]]: """Returns the partial tool calls for the 0th choice message. @@ -297,3 +317,22 @@ def tool_calls(self) -> Optional[list[ChoiceDeltaToolCall]]: `list[ChoiceDeltaToolCall]` will be None indicating end of stream. """ return self.delta.tool_calls + + @property + def usage(self) -> Optional[CompletionUsage]: + """Returns the usage of the chat completion.""" + return self.chunk.usage + + @property + def input_tokens(self) -> Optional[int]: + """Returns the number of input tokens.""" + if self.usage: + return self.usage.prompt_tokens + return None + + @property + def output_tokens(self) -> Optional[int]: + """Returns the number of output tokens.""" + if self.usage: + return self.usage.completion_tokens + return None diff --git a/mirascope/mistral/types.py b/mirascope/mistral/types.py index 06096eb9..e1ef85b3 100644 --- a/mirascope/mistral/types.py +++ b/mirascope/mistral/types.py @@ -91,15 +91,23 @@ def content(self) -> str: # return the first item in the list return content if isinstance(content, str) else content[0] + @property + def model(self) -> str: + """Returns the name of the response model.""" + return self.response.model + @property def id(self) -> str: """Returns the id of the response.""" return self.response.id @property - def finish_reasons(self) -> list[Optional[str]]: + def finish_reasons(self) -> list[str]: """Returns the finish reasons of the response.""" - return [choice.finish_reason for choice in self.choices] + return [ + choice.finish_reason if choice.finish_reason else "" + for choice in self.choices + ] @property def tool_calls(self) -> Optional[list[ToolCall]]: @@ -223,7 +231,40 @@ def content(self) -> str: """Returns the content of the delta.""" return self.delta.content if self.delta.content is not None else "" + @property + def model(self) -> str: + """Returns the name of the response model.""" + return self.chunk.model + + @property + def id(self) -> str: + """Returns the id of the response.""" + return self.chunk.id + + @property + def finish_reasons(self) -> list[str]: + """Returns the finish reasons of the response.""" + return [ + choice.finish_reason if choice.finish_reason else "" + for choice in self.choices + ] + @property def tool_calls(self) -> Optional[list[ToolCall]]: """Returns the partial tool calls for the 0th choice message.""" return self.delta.tool_calls + + @property + def usage(self) -> UsageInfo: + """Returns the usage of the chat completion.""" + return self.chunk.usage + + @property + def input_tokens(self) -> int: + """Returns the number of input tokens.""" + return self.usage.prompt_tokens + + @property + def output_tokens(self) -> Optional[int]: + """Returns the number of output tokens.""" + return self.usage.completion_tokens diff --git a/mirascope/openai/types.py b/mirascope/openai/types.py index e7fb128f..f8da5b33 100644 --- a/mirascope/openai/types.py +++ b/mirascope/openai/types.py @@ -124,6 +124,11 @@ def content(self) -> str: """Returns the content of the chat completion for the 0th choice.""" return self.message.content if self.message.content is not None else "" + @property + def model(self) -> str: + """Returns the name of the response model.""" + return self.response.model + @property def id(self) -> str: """Returns the id of the response.""" @@ -132,7 +137,7 @@ def id(self) -> str: @property def finish_reasons(self) -> list[str]: """Returns the finish reasons of the response.""" - return [choice.finish_reason for choice in self.response.choices] + return [str(choice.finish_reason) for choice in self.response.choices] @property def tool_calls(self) -> Optional[list[ChatCompletionMessageToolCall]]: @@ -292,6 +297,11 @@ def content(self) -> str: self.delta.content if self.delta is not None and self.delta.content else "" ) + @property + def model(self) -> str: + """Returns the name of the response model.""" + return self.chunk.model + @property def id(self) -> str: """Returns the id of the response.""" @@ -300,7 +310,7 @@ def id(self) -> str: @property def finish_reasons(self) -> list[str]: """Returns the finish reasons of the response.""" - return [choice.finish_reason for choice in self.chunk.choices] + return [str(choice.finish_reason) for choice in self.chunk.choices] @property def tool_calls(self) -> Optional[list[ChoiceDeltaToolCall]]: diff --git a/mirascope/otel/otel.py b/mirascope/otel/otel.py index 9b77a321..f340e35e 100644 --- a/mirascope/otel/otel.py +++ b/mirascope/otel/otel.py @@ -8,7 +8,7 @@ ) from typing import Any, Callable, Optional, Sequence, Union, overload -from opentelemetry.sdk.trace import Span, SpanProcessor, TracerProvider +from opentelemetry.sdk.trace import SpanProcessor, TracerProvider from opentelemetry.sdk.trace.export import ( ConsoleSpanExporter, SimpleSpanProcessor, @@ -20,6 +20,8 @@ get_tracer_provider, set_tracer_provider, ) +from opentelemetry.trace.span import Span +from opentelemetry.util.types import Attributes, AttributeValue from pydantic import BaseModel from typing_extensions import LiteralString @@ -45,28 +47,31 @@ ) -def _set_response_attributes( - span: Span, response: BaseCallResponse, messages: list[dict[str, Any]] -) -> None: - """Sets the response attributes""" - response_attributes = { - "gen_ai.response.id": response.id, - "gen_ai.response.model": response.content, - "gen_ai.response.finish_reasons": response.finish_reasons, - "gen_ai.usage.completion_tokens": response.output_tokens, - "gen_ai.usage.prompt_tokens": response.input_tokens, - } - span.set_attributes(response_attributes) - event_attributes = {"gen_ai.completion": json.dumps(messages)} - span.add_event( - "gen_ai.content.completion", - attributes=event_attributes, - ) - - def mirascope_otel(cls) -> Callable: tracer = get_tracer("otel") + def _set_response_attributes( + span: Span, response: BaseCallResponse, messages: list[dict[str, Any]] + ) -> None: + """Sets the response attributes""" + response_attributes: dict[str, AttributeValue] = {} + if model := response.model: + response_attributes["gen_ai.response.model"] = model + if id := response.id: + response_attributes["gen_ai.response.id"] = id + if finish_reasons := response.finish_reasons: + response_attributes["gen_ai.response.finish_reasons"] = finish_reasons + if output_tokens := response.output_tokens: + response_attributes["gen_ai.usage.completion_tokens"] = output_tokens + if input_tokens := response.input_tokens: + response_attributes["gen_ai.usage.prompt_tokens"] = input_tokens + span.set_attributes(response_attributes) + event_attributes: Attributes = {"gen_ai.completion": json.dumps(messages)} + span.add_event( + "gen_ai.content.completion", + attributes=event_attributes, + ) + def _set_span_data( span: Span, suffix: str, @@ -74,8 +79,8 @@ def _set_span_data( model_name: Optional[str], args: tuple[Any], ) -> None: - prompt = args[0] - if suffix == "gemini": + prompt = args[0] if len(args) > 0 else None + if prompt and suffix == "gemini": prompt = { "messages": [ {"role": message["role"], "content": message["parts"][0]} @@ -91,15 +96,20 @@ def _set_span_data( max_tokens = getattr(call_params, "max_tokens", None) temperature = getattr(call_params, "temperature", None) top_p = getattr(call_params, "top_p", None) - attributes = { + attributes: dict[str, AttributeValue] = { "async": is_async, - "gen_ai.request.model": model_name, "gen_ai.system": suffix, - "gen_ai.request.max_tokens": max_tokens, - "gen_ai.request.temperature": temperature, - "gen_ai.request.top_p": top_p, } - events = { + if model_name: + attributes["gen_ai.request.model"] = model_name + if max_tokens: + attributes["gen_ai.request.max_tokens"] = max_tokens + if temperature: + attributes["gen_ai.request.temperature"] = temperature + if top_p: + attributes["gen_ai.request.top_p"] = top_p + + events: Attributes = { "gen_ai.prompt": json.dumps(prompt), } span.set_attributes(attributes) @@ -132,31 +142,46 @@ def record_streaming( ): """OTel record_streaming with Mirascope providers""" content: list[str] = [] + response_attributes: dict[str, AttributeValue] = {} def record_chunk( - chunk: ChunkT, response_chunk_type: type[BaseCallResponseChunk] + raw_chunk: ChunkT, response_chunk_type: type[BaseCallResponseChunk] ) -> Any: """Handles all provider chunk_types instead of only OpenAI""" - chunk_content = content_from_stream(chunk, response_chunk_type) + chunk = content_from_stream(raw_chunk, response_chunk_type) + if model := chunk.model: + response_attributes["gen_ai.response.model"] = model + if id := chunk.id: + response_attributes["gen_ai.response.id"] = id + if finish_reasons := chunk.finish_reasons: + response_attributes["gen_ai.response.finish_reasons"] = finish_reasons + if output_tokens := chunk.output_tokens: + response_attributes["gen_ai.usage.completion_tokens"] = output_tokens + if input_tokens := chunk.input_tokens: + response_attributes["gen_ai.usage.prompt_tokens"] = input_tokens + chunk_content = chunk.content if chunk_content is not None: content.append(chunk_content) try: yield record_chunk finally: - attributes = { - "response_data": { - "combined_chunk_content": "".join(content), - "chunk_count": len(content), - }, + span.set_attributes(response_attributes) + event_attributes: Attributes = { + "gen_ai.completion": json.dumps( + {"role": "assistant", "content": "".join(content)} + ) } - span.set_attributes(attributes) + span.add_event( + "gen_ai.content.completion", + attributes=event_attributes, + ) - def _extract_chunk_content( + def _extract_chunk( chunk: ChunkT, response_chunk_type: type[BaseCallResponseChunk] ) -> str: """Extracts the content from a chunk.""" - return response_chunk_type(chunk=chunk).content + return response_chunk_type(chunk=chunk) def mirascope_otel_decorator( fn: Callable, @@ -230,7 +255,7 @@ def wrapper_generator(*args, **kwargs): f"{suffix}.{fn.__name__} with {model}" ) as span: _set_span_data(span, suffix, is_async, model, args) - with record_streaming(span, _extract_chunk_content) as record_chunk: + with record_streaming(span, _extract_chunk) as record_chunk: stream = fn(*args, **kwargs) if isinstance(stream, AbstractContextManager): with stream as s: @@ -248,7 +273,7 @@ async def wrapper_generator_async(*args, **kwargs): f"{suffix}.{fn.__name__} with {model}" ) as span: _set_span_data(span, suffix, is_async, model, args) - with record_streaming(span, _extract_chunk_content) as record_chunk: + with record_streaming(span, _extract_chunk) as record_chunk: stream = fn(*args, **kwargs) if inspect.iscoroutine(stream): stream = await stream diff --git a/tests/anthropic/conftest.py b/tests/anthropic/conftest.py index 234d3677..35fa6506 100644 --- a/tests/anthropic/conftest.py +++ b/tests/anthropic/conftest.py @@ -60,6 +60,7 @@ def fixture_anthropic_message() -> Message: role="assistant", type="message", usage=Usage(input_tokens=0, output_tokens=0), + stop_reason="end_turn", ) diff --git a/tests/anthropic/test_types.py b/tests/anthropic/test_types.py index 1fd7ea86..283fa39a 100644 --- a/tests/anthropic/test_types.py +++ b/tests/anthropic/test_types.py @@ -8,7 +8,9 @@ ContentBlockStartEvent, ContentBlockStopEvent, Message, + MessageStartEvent, TextBlock, + Usage, ) from mirascope.anthropic.tools import AnthropicTool @@ -26,6 +28,9 @@ def test_anthropic_call_response(fixture_anthropic_message: Message): assert response.usage is not None assert response.output_tokens is not None assert response.input_tokens is not None + assert response.model == fixture_anthropic_message.model + assert response.id == fixture_anthropic_message.id + assert response.finish_reasons == [fixture_anthropic_message.stop_reason] assert response.dump() == { "start_time": 0, "end_time": 1, @@ -89,3 +94,24 @@ def test_anthropic_call_response_chunk( ) assert chunk.content == "" assert chunk.type == "content_block_stop" + + chunk = AnthropicCallResponseChunk( + chunk=MessageStartEvent( + message=Message( + id="test_id", + model="test_model", + role="assistant", + type="message", + content=[TextBlock(text="test", type="text")], + stop_reason="end_turn", + usage=Usage(input_tokens=1, output_tokens=1), + ), + type="message_start", + ) + ) + assert chunk.usage is not None + assert chunk.output_tokens == 1 + assert chunk.input_tokens == 1 + assert chunk.model == "test_model" + assert chunk.id == "test_id" + assert chunk.finish_reasons == ["end_turn"] diff --git a/tests/cohere/test_types.py b/tests/cohere/test_types.py index dbd3df4e..70b50725 100644 --- a/tests/cohere/test_types.py +++ b/tests/cohere/test_types.py @@ -11,6 +11,7 @@ StreamedChatResponse_SearchQueriesGeneration, StreamedChatResponse_SearchResults, StreamedChatResponse_StreamEnd, + StreamedChatResponse_StreamStart, StreamedChatResponse_TextGeneration, StreamedChatResponse_ToolCallsGeneration, ToolCall, @@ -124,10 +125,16 @@ def test_cohere_call_response_chunk_properties( assert call_chunk.event_type == "text-generation" assert call_chunk.content == "Test response" assert call_chunk.response is None + call_chunk.chunk = StreamedChatResponse_StreamStart( + generation_id="test_id", + event_type="stream-start", + **fixture_non_streamed_response.dict(exclude={"event_type"}), + ) + assert call_chunk.id == "test_id" call_chunk.chunk = StreamedChatResponse_SearchQueriesGeneration( event_type="search-queries-generation", - **chunk.dict(exclude={"event_type"}), + **fixture_non_streamed_response.dict(exclude={"event_type"}), ) assert call_chunk.search_queries == [fixture_chat_search_query] assert call_chunk.search_results is None @@ -178,3 +185,7 @@ def test_cohere_call_response_chunk_properties( assert call_chunk.documents is None assert call_chunk.citations is None assert call_chunk.content == "" + assert call_chunk.finish_reasons == ["COMPLETE"] + assert call_chunk.usage is not None + assert call_chunk.input_tokens is not None + assert call_chunk.output_tokens is not None diff --git a/tests/conftest.py b/tests/conftest.py index 0552437c..673f0b13 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -436,9 +436,10 @@ def fixture_generate_content_response(): GenerateContentResponse( candidates=[ Candidate( + finish_reason=1, content=Content( parts=[Part(text="Who is the author?")], role="model" - ) + ), ) ] ) diff --git a/tests/gemini/conftest.py b/tests/gemini/conftest.py index b08e6acf..508177e7 100644 --- a/tests/gemini/conftest.py +++ b/tests/gemini/conftest.py @@ -92,18 +92,26 @@ def fixture_expected_book_tool_instance() -> BookTool: @pytest.fixture() -def fixture_generate_content_chunks(): +def fixture_generate_content_chunks() -> GenerateContentResponse: """Returns a list of `GenerateContentResponse` chunks.""" return GenerateContentResponseType.from_iterator( [ GenerateContentResponse( candidates=[ - Candidate(content=Content(parts=[{"text": "first"}], role="model")) + Candidate( + finish_reason=1, + index=0, + content=Content(parts=[{"text": "first"}], role="model"), + ) ] ), GenerateContentResponse( candidates=[ - Candidate(content=Content(parts=[{"text": "second"}], role="model")) + Candidate( + finish_reason=1, + index=0, + content=Content(parts=[{"text": "second"}], role="model"), + ) ] ), ] diff --git a/tests/gemini/test_types.py b/tests/gemini/test_types.py index 1d6d9907..a4584612 100644 --- a/tests/gemini/test_types.py +++ b/tests/gemini/test_types.py @@ -6,7 +6,7 @@ from google.generativeai.types import GenerateContentResponse # type: ignore from mirascope.gemini.tools import GeminiTool -from mirascope.gemini.types import GeminiCallResponse +from mirascope.gemini.types import GeminiCallResponse, GeminiCallResponseChunk def test_gemini_call_response( @@ -102,3 +102,21 @@ def test_gemini_call_response_with_no_tools( assert response.tools is None assert response.tool is None + + +def test_gemini_call_response_chunk( + fixture_generate_content_response_with_tools: GenerateContentResponse, +) -> None: + """Tests the `GeminiCallResponseChunk` class.""" + response = GeminiCallResponseChunk( + chunk=fixture_generate_content_response_with_tools, + start_time=0, + end_time=0, + ) + assert response.chunk is not None + assert response.id is None + assert response.finish_reasons == ["STOP"] + assert response.usage is None + assert response.model is None + assert response.input_tokens is None + assert response.output_tokens is None diff --git a/tests/groq/conftest.py b/tests/groq/conftest.py index 60ec8d7e..7ff2ee58 100644 --- a/tests/groq/conftest.py +++ b/tests/groq/conftest.py @@ -250,6 +250,7 @@ def fixture_chat_completion_stream_response_with_tools() -> list[ChatCompletionC object="chat.completion.chunk", system_fingerprint="", x_groq=None, + usage=CompletionUsage(prompt_tokens=1, completion_tokens=1, total_tokens=2), ), ChatCompletionChunk( id="test", @@ -270,6 +271,7 @@ def fixture_chat_completion_stream_response_with_tools() -> list[ChatCompletionC object="chat.completion.chunk", system_fingerprint="", x_groq=None, + usage=None, ), ] diff --git a/tests/groq/test_types.py b/tests/groq/test_types.py index 37205562..84425f08 100644 --- a/tests/groq/test_types.py +++ b/tests/groq/test_types.py @@ -25,6 +25,11 @@ def test_groq_call_response( assert response.content == "test content" assert response.tools is None assert response.tool is None + assert response.model == fixture_chat_completion_response.model + assert response.id == fixture_chat_completion_response.id + assert response.finish_reasons == [ + fixture_chat_completion_response.choices[0].finish_reason + ] assert response.usage is not None assert response.input_tokens is not None assert response.output_tokens is not None @@ -162,3 +167,19 @@ def test_groq_stream_response_with_tools( .choices[0] .delta.tool_calls ) + assert response.model == fixture_chat_completion_stream_response_with_tools[0].model + assert response.id == fixture_chat_completion_stream_response_with_tools[0].id + assert response.finish_reasons == [ + fixture_chat_completion_stream_response_with_tools[0].choices[0].finish_reason + ] + assert response.usage is not None + assert response.input_tokens == 1 + assert response.output_tokens == 1 + + response = GroqCallResponseChunk( + chunk=fixture_chat_completion_stream_response_with_tools[1], + tool_types=[fixture_book_tool], + ) + assert response.usage is None + assert response.input_tokens is None + assert response.output_tokens is None diff --git a/tests/mistral/conftest.py b/tests/mistral/conftest.py index 429d98f3..9b31a939 100644 --- a/tests/mistral/conftest.py +++ b/tests/mistral/conftest.py @@ -40,7 +40,7 @@ def fixture_chat_completion_response( model="open-mixtral-8x7b", choices=[ ChatCompletionResponseChoice( - index=0, message=fixture_chat_message, finish_reason=None + index=0, message=fixture_chat_message, finish_reason=FinishReason.stop ) ], usage=UsageInfo(prompt_tokens=1, total_tokens=2, completion_tokens=1), @@ -173,7 +173,7 @@ def fixture_chat_completion_stream_response_with_tools() -> ( role="assistant", tool_calls=[tool_call], ), - finish_reason=None, + finish_reason=FinishReason.tool_calls, ) ], usage=UsageInfo(prompt_tokens=1, total_tokens=2, completion_tokens=1), @@ -188,7 +188,7 @@ def fixture_chat_completion_stream_response_with_tools() -> ( role="assistant", tool_calls=[tool_call], ), - finish_reason=None, + finish_reason=FinishReason.tool_calls, ) ], usage=UsageInfo(prompt_tokens=1, total_tokens=2, completion_tokens=1), diff --git a/tests/mistral/test_types.py b/tests/mistral/test_types.py index aa6bdd09..1cde3e3e 100644 --- a/tests/mistral/test_types.py +++ b/tests/mistral/test_types.py @@ -27,6 +27,11 @@ def test_mistral_call_response( assert response.content == "test content" assert response.tools is None assert response.tool is None + assert response.model == fixture_chat_completion_response.model + assert response.id == fixture_chat_completion_response.id + assert response.finish_reasons == [ + fixture_chat_completion_response.choices[0].finish_reason + ] assert response.usage is not None assert response.input_tokens is not None assert response.output_tokens is not None @@ -137,3 +142,11 @@ def test_mistral_stream_response_with_tools( .choices[0] .delta.tool_calls ) + assert response.model == fixture_chat_completion_stream_response_with_tools[0].model + assert response.id == fixture_chat_completion_stream_response_with_tools[0].id + assert response.finish_reasons == [ + fixture_chat_completion_stream_response_with_tools[0].choices[0].finish_reason + ] + assert response.usage is not None + assert response.input_tokens is not None + assert response.output_tokens is not None diff --git a/tests/openai/test_types.py b/tests/openai/test_types.py index 27c58ffd..7e52cc39 100644 --- a/tests/openai/test_types.py +++ b/tests/openai/test_types.py @@ -32,6 +32,9 @@ def test_openai_call_response(fixture_chat_completion: ChatCompletion): assert response.choice == choices[0] assert response.message == choices[0].message assert response.content == choices[0].message.content + assert response.model == fixture_chat_completion.model + assert response.finish_reasons == ["stop", "stop"] + assert response.id == fixture_chat_completion.id assert response.tools is None assert response.tool is None @@ -134,6 +137,9 @@ def test_openai_chat_completion_chunk( assert openai_chat_completion_chunk.choice == choices[0] assert openai_chat_completion_chunk.delta == choices[0].delta assert openai_chat_completion_chunk.content == choices[0].delta.content + assert openai_chat_completion_chunk.model == fixture_chat_completion_chunk.model + assert openai_chat_completion_chunk.id == fixture_chat_completion_chunk.id + assert openai_chat_completion_chunk.finish_reasons == ["stop"] def test_openai_chat_completion_last_chunk( diff --git a/tests/otel/test_otel.py b/tests/otel/test_otel.py index 2235cf84..c8e52468 100644 --- a/tests/otel/test_otel.py +++ b/tests/otel/test_otel.py @@ -62,7 +62,9 @@ class MyCall(OpenAICall): class MyNestedCall(MyCall): prompt_template = "test" - call_params = OpenAICallParams(model="gpt-3.5-turbo") + call_params = OpenAICallParams( + model="gpt-3.5-turbo", temperature=0.1, top_p=0.9 + ) configuration = BaseConfig(llm_ops=[], client_wrappers=[]) MyNestedCall().call() @@ -337,7 +339,7 @@ def test_value_error_on_mirascope_otel(): def foo(): ... # pragma: no cover - mirascope_otel()(foo, "test") + mirascope_otel(BaseModel)(foo, "test") @patch("cohere.Client.chat", new_callable=MagicMock) From 124a6c400f81aaf7dc1ecffb326019424513fc9c Mon Sep 17 00:00:00 2001 From: Brendan Date: Thu, 6 Jun 2024 16:37:40 -0700 Subject: [PATCH 3/3] fixed set_attribute errors with otel --- mirascope/anthropic/types.py | 4 ++-- mirascope/base/types.py | 27 +++++++++++++++++++++++++++ mirascope/cohere/types.py | 4 ++-- mirascope/mistral/types.py | 12 ++++++++---- mirascope/otel/otel.py | 29 +++++++++++++++++++++++++---- tests/gemini/test_types.py | 4 +--- tests/mistral/conftest.py | 2 +- tests/mistral/test_types.py | 9 +++++++++ 8 files changed, 75 insertions(+), 16 deletions(-) diff --git a/mirascope/anthropic/types.py b/mirascope/anthropic/types.py index 24f55660..bd59af3a 100644 --- a/mirascope/anthropic/types.py +++ b/mirascope/anthropic/types.py @@ -253,7 +253,7 @@ def content(self) -> str: return "" @property - def model(self) -> str: + def model(self) -> Optional[str]: """Returns the name of the response model.""" if isinstance(self.chunk, MessageStartEvent): return self.chunk.message.model @@ -270,7 +270,7 @@ def id(self) -> Optional[str]: def finish_reasons(self) -> Optional[list[str]]: """Returns the finish reason of the response.""" if isinstance(self.chunk, MessageStartEvent): - return [self.chunk.message.stop_reason] + return [str(self.chunk.message.stop_reason)] return None @property diff --git a/mirascope/base/types.py b/mirascope/base/types.py index df752cae..a225f2d8 100644 --- a/mirascope/base/types.py +++ b/mirascope/base/types.py @@ -260,3 +260,30 @@ def finish_reasons(self) -> Union[None, list[str]]: If there is no finish reason, this method must return None. """ ... # pragma: no cover + + @property + @abstractmethod + def usage(self) -> Any: + """Should return the usage of the response. + + If there is no usage, this method must return None. + """ + ... # pragma: no cover + + @property + @abstractmethod + def input_tokens(self) -> Optional[Union[int, float]]: + """Should return the number of input tokens. + + If there is no input_tokens, this method must return None. + """ + ... # pragma: no cover + + @property + @abstractmethod + def output_tokens(self) -> Optional[Union[int, float]]: + """Should return the number of output tokens. + + If there is no output_tokens, this method must return None. + """ + ... # pragma: no cover diff --git a/mirascope/cohere/types.py b/mirascope/cohere/types.py index b4f7e6d2..54a94ab4 100644 --- a/mirascope/cohere/types.py +++ b/mirascope/cohere/types.py @@ -123,7 +123,7 @@ def content(self) -> str: return self.response.text @property - def model(self) -> str: + def model(self) -> Optional[str]: """Returns the name of the response model. Cohere does not return model, so we return None @@ -318,7 +318,7 @@ def citations(self) -> Optional[list[ChatCitation]]: return None @property - def model(self) -> str: + def model(self) -> Optional[str]: """Returns the name of the response model. Cohere does not return model, so we return None diff --git a/mirascope/mistral/types.py b/mirascope/mistral/types.py index e1ef85b3..54967a5e 100644 --- a/mirascope/mistral/types.py +++ b/mirascope/mistral/types.py @@ -255,16 +255,20 @@ def tool_calls(self) -> Optional[list[ToolCall]]: return self.delta.tool_calls @property - def usage(self) -> UsageInfo: + def usage(self) -> Optional[UsageInfo]: """Returns the usage of the chat completion.""" return self.chunk.usage @property - def input_tokens(self) -> int: + def input_tokens(self) -> Optional[int]: """Returns the number of input tokens.""" - return self.usage.prompt_tokens + if self.usage: + return self.usage.prompt_tokens + return None @property def output_tokens(self) -> Optional[int]: """Returns the number of output tokens.""" - return self.usage.completion_tokens + if self.usage: + return self.usage.completion_tokens + return None diff --git a/mirascope/otel/otel.py b/mirascope/otel/otel.py index f340e35e..0a8ba912 100644 --- a/mirascope/otel/otel.py +++ b/mirascope/otel/otel.py @@ -137,7 +137,7 @@ def _mirascope_llm_span( def record_streaming( span: Span, content_from_stream: Callable[ - [ChunkT, type[BaseCallResponseChunk]], Union[str, None] + [ChunkT, type[BaseCallResponseChunk]], BaseCallResponseChunk ], ): """OTel record_streaming with Mirascope providers""" @@ -148,7 +148,9 @@ def record_chunk( raw_chunk: ChunkT, response_chunk_type: type[BaseCallResponseChunk] ) -> Any: """Handles all provider chunk_types instead of only OpenAI""" - chunk = content_from_stream(raw_chunk, response_chunk_type) + chunk: BaseCallResponseChunk = content_from_stream( + raw_chunk, response_chunk_type + ) if model := chunk.model: response_attributes["gen_ai.response.model"] = model if id := chunk.id: @@ -179,7 +181,7 @@ def record_chunk( def _extract_chunk( chunk: ChunkT, response_chunk_type: type[BaseCallResponseChunk] - ) -> str: + ) -> BaseCallResponseChunk: """Extracts the content from a chunk.""" return response_chunk_type(chunk=chunk) @@ -300,6 +302,11 @@ async def wrapper_generator_async(*args, **kwargs): return mirascope_otel_decorator +def custom_encoder(obj) -> str: + """Custom encoder for the OpenTelemetry span""" + return obj.__name__ + + @contextmanager def handle_before_call(self: BaseModel, fn: Callable, **kwargs): """Handles before call""" @@ -307,10 +314,24 @@ def handle_before_call(self: BaseModel, fn: Callable, **kwargs): class_vars = get_class_vars(self) inputs = self.model_dump() tracer = get_tracer("otel") + attributes: dict[str, AttributeValue] = {**kwargs, **class_vars, **inputs} + if hasattr(self, "call_params"): + attributes["call_params"] = json.dumps( + self.call_params.model_dump(), default=custom_encoder + ) + if hasattr(self, "configuration"): + configuration = self.configuration.model_dump() + attributes["configuration"] = json.dumps(configuration, default=custom_encoder) + + if hasattr(self, "base_url"): + attributes["base_url"] = self.base_url if self.base_url else "" + + if hasattr(self, "extract_schema"): + attributes["extract_schema"] = self.extract_schema.__name__ with tracer.start_as_current_span( f"{self.__class__.__name__}.{fn.__name__}" ) as span: - span.set_attributes({**kwargs, **class_vars, **inputs}) + span.set_attributes(attributes) yield span diff --git a/tests/gemini/test_types.py b/tests/gemini/test_types.py index a4584612..d3a7b072 100644 --- a/tests/gemini/test_types.py +++ b/tests/gemini/test_types.py @@ -109,9 +109,7 @@ def test_gemini_call_response_chunk( ) -> None: """Tests the `GeminiCallResponseChunk` class.""" response = GeminiCallResponseChunk( - chunk=fixture_generate_content_response_with_tools, - start_time=0, - end_time=0, + chunk=fixture_generate_content_response_with_tools ) assert response.chunk is not None assert response.id is None diff --git a/tests/mistral/conftest.py b/tests/mistral/conftest.py index 9b31a939..2386c44b 100644 --- a/tests/mistral/conftest.py +++ b/tests/mistral/conftest.py @@ -191,6 +191,6 @@ def fixture_chat_completion_stream_response_with_tools() -> ( finish_reason=FinishReason.tool_calls, ) ], - usage=UsageInfo(prompt_tokens=1, total_tokens=2, completion_tokens=1), + usage=None, ), ] diff --git a/tests/mistral/test_types.py b/tests/mistral/test_types.py index 1cde3e3e..6fc4f066 100644 --- a/tests/mistral/test_types.py +++ b/tests/mistral/test_types.py @@ -150,3 +150,12 @@ def test_mistral_stream_response_with_tools( assert response.usage is not None assert response.input_tokens is not None assert response.output_tokens is not None + + response = MistralCallResponseChunk( + chunk=fixture_chat_completion_stream_response_with_tools[1], + tool_types=[fixture_book_tool], + ) + + assert response.usage is None + assert response.input_tokens is None + assert response.output_tokens is None