Skip to content

Commit

Permalink
improving code
Browse files Browse the repository at this point in the history
  • Loading branch information
Redna committed Jun 13, 2024
1 parent e2d6f33 commit 70afb75
Showing 1 changed file with 2 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -125,14 +125,12 @@ def trace(self, operation_name: str, tags: Optional[Dict[str, Any]] = None) -> I
# Haystack returns one meta dict for each message, but the 'usage' value
# is always the same, let's just pick the first item
m = meta[0]
usage = m.get("usage") if m.get("usage") else None
span._span.update(usage=usage, model=m.get("model"))
span._span.update(usage=m.get("usage") or None, model=m.get("model"))
elif tags.get("haystack.component.type") == "OpenAIChatGenerator":
replies = span._data.get("haystack.component.output", {}).get("replies")
if replies:
meta = replies[0].meta
usage = meta.get("usage") if meta.get("usage") else None
span._span.update(usage=usage, model=meta.get("model"))
span._span.update(usage=meta.get("usage") or None, model=meta.get("model"))

pipeline_input = tags.get("haystack.pipeline.input_data", None)
if pipeline_input:
Expand Down

0 comments on commit 70afb75

Please sign in to comment.