diff --git a/literalai/exporter.py b/literalai/exporter.py index 28877e3..2096116 100644 --- a/literalai/exporter.py +++ b/literalai/exporter.py @@ -114,17 +114,17 @@ def _create_step_from_span(self, span: ReadableSpan) -> Step: else None ), ) + messages = self._extract_messages(cast(Dict, attributes)) if is_chat else [] + + message_completions = self._extract_messages(cast(Dict, attributes), "gen_ai.completion.") if is_chat else [] + + message_completion = message_completions[-1] if message_completions else None + previous_messages = messages + message_completions[:-1] if message_completions else messages generation_content = { "duration": duration, - "messages": ( - self._extract_messages(cast(Dict, attributes)) if is_chat else None - ), - "message_completion": ( - self._extract_messages(cast(Dict, attributes), "gen_ai.completion.")[0] - if is_chat - else None - ), + "messages": previous_messages, + "message_completion": message_completion, "prompt": attributes.get("gen_ai.prompt.0.user"), "promptId": prompt.get("id") if prompt else None, "completion": attributes.get("gen_ai.completion.0.content"), diff --git a/literalai/version.py b/literalai/version.py index 7096515..1e9961d 100644 --- a/literalai/version.py +++ b/literalai/version.py @@ -1 +1 @@ -__version__ = "0.1.104" +__version__ = "0.1.105" diff --git a/setup.py b/setup.py index 880f31f..3b32860 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ setup( name="literalai", - version="0.1.104", # update version in literalai/version.py + version="0.1.105", # update version in literalai/version.py description="An SDK for observability in Python applications", long_description=open("README.md").read(), long_description_content_type="text/markdown",