[truecallerabreham]

Summary

This PR fixes multiple unhandled exceptions in the Python SDK's LiteLLM integration (sdks/python/agenta/sdk/litellm/litellm.py) and resolves broken lazy loading in assets.py:

1. TypeError: 'NoneType' object is not a mapping: In log_pre_api_call, **kwargs.get("optional_params") crashed when optional_params was absent or None. It now defaults safely to {}.
2. KeyError: 'messages' for non-chat operations: In log_pre_api_call, accessing kwargs["messages"] crashed on embedding or completion calls where inputs are passed under "prompt" or "input". Input extraction now safely falls back across kwargs.get("messages"), messages, kwargs.get("prompt"), and kwargs.get("input").
3. Unhandled KeyError on span lookups: In log_stream_event, log_success_event, log_failure_event, and their async variants, self.span[litellm_call_id] raised an unhandled KeyError if a span was untracked or log_pre_api_call failed, bypassing the intended if not span: guard. Switched to self.span.get(litellm_call_id).
4. Unhandled KeyError on span deletion: del self.span[litellm_call_id] crashed if already cleaned up or untracked. Replaced with self.span.pop(litellm_call_id, None).
5. Unhandled KeyError: 'exception': span.record_exception(kwargs["exception"]) crashed if kwargs did not contain "exception". Added safe key and None check.
6. Preserved lazy loading in assets.py: Top-level from litellm import cost_calculator broke lazy loading across the SDK and caused any import of agenta to crash when litellm was not installed. Wrapped in try/except Exception: cost_calculator = None.

Testing

Verified locally

• Ran full test execution with pytest against the Python SDK unit test suite.
• Verified that all 10 edge cases reproduce the crashes prior to the fix and pass cleanly after the fix.
• Ran all existing SDK unit tests (test_message_types.py, test_preinit_warnings.py, test_redaction.py, test_redaction_seed.py, test_structured_rendering.py) — total 79 tests passing with 0 regressions.

Added or updated tests

• Created sdks/python/oss/tests/pytest/unit/test_litellm_handler.py containing 10 automated test cases testing sync and async callback safety, missing/None optional params, embedding calls without messages, missing exceptions, and unknown call IDs.

QA follow-up

• N/A

Demo

Checklist

• I have included a video or screen recording for UI changes, or marked Demo as N/A
• Relevant tests pass locally
• Relevant linting and formatting pass locally
• I have signed the CLA, or I will sign it when the bot prompts me

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository YAML (base), Organization UI (inherited)

Review profile: CHILL

Plan: Advanced

Run ID: 55da538f-94bc-4a25-962f-f740d7083588

📥 Commits

Reviewing files that changed from the base of the PR and between d57f0a2 and 0579d82.

⛔ Files ignored due to path filters (1)

• sdks/python/oss/tests/pytest/unit/test_litellm_results.svg is excluded by !**/*.svg

📒 Files selected for processing (2)

• sdks/python/agenta/sdk/litellm/litellm.py
• sdks/python/oss/tests/pytest/unit/test_litellm_handler.py

🚧 Files skipped from review as they are similar to previous changes (2)

• sdks/python/oss/tests/pytest/unit/test_litellm_handler.py
• sdks/python/agenta/sdk/litellm/litellm.py

Included review availability: Your plan provides up to 8 included reviews per hour; 6 remain after this review.

---

📝 Summary

Summary by CodeRabbit

• Bug Fixes

  • Improved reliability of LiteLLM request tracing across synchronous and asynchronous calls.
  • Prevented errors when optional request parameters, messages, or exception details are missing.
  • Preserved valid empty prompt values during request handling.
  • Ensured completed and failed requests are cleaned up safely.
  • LiteLLM is now optional when collecting model cost metadata, allowing other functionality to continue when it is unavailable.

• Tests

  • Added coverage for embedding calls, positional messages, missing parameters, and unknown request IDs.

Walkthrough

The LiteLLM handler now handles missing spans, optional exceptions, and multiple input formats in synchronous and asynchronous callbacks. LiteLLM cost calculation imports are optional. Unit tests cover these cases.

Changes

LiteLLM robustness

Layer / File(s)	Summary
Callback state and input handling sdks/python/agenta/sdk/litellm/litellm.py, sdks/python/oss/tests/pytest/unit/test_litellm_handler.py	Callbacks use safe span lookup and cleanup. Failure callbacks record exceptions only when provided. Pre-call handling supports fallback input sources and missing optional_params. Tests cover synchronous and asynchronous callbacks.
Optional LiteLLM cost loading sdks/python/agenta/sdk/utils/assets.py	The module catches LiteLLM import failures and sets cost_calculator to None.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~20 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 60.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 22 functions across 3 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main changes: defensive safeguards for the LiteLLM handler and restored lazy loading for assets.
Description check	✅ Passed	The description directly explains the LiteLLM error fixes, lazy loading change, tests, and reported regression status.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository YAML (base), Organization UI (inherited)

Review profile: CHILL

Plan: Advanced

Run ID: 4c1f29b5-461c-4fb6-9ebd-fd32a6147270

📥 Commits

Reviewing files that changed from the base of the PR and between 21fc6af and d57f0a2.

⛔ Files ignored due to path filters (1)

• sdks/python/oss/tests/pytest/unit/test_litellm_results.svg is excluded by !**/*.svg

📒 Files selected for processing (3)

• sdks/python/agenta/sdk/litellm/litellm.py
• sdks/python/agenta/sdk/utils/assets.py
• sdks/python/oss/tests/pytest/unit/test_litellm_handler.py

Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review.