Expose unified model calls to extensions
Decorate unified_call and unified_turn so community plugins can observe the real logical model-call boundary. Add focused coverage for the generated start and end extension points.
Alessandro committed
Jul 17, 2026 at 17:00 UTC
dc7004462efc9de533077edf6874a85f81bcd398
2 files changed
+36
models.py
+2
@@ -524,6 +524,7 @@ class LiteLLMChatWrapper(SimpleChatModel):
524
message=AIMessageChunk(content=output["response_delta"])
525
)
526
527
+ @extensible
528
async def unified_call(
529
self,
530
system_message="",
@@ -637,6 +638,7 @@ class LiteLLMChatWrapper(SimpleChatModel):
638
attempt += 1
639
await asyncio.sleep(retry_delay_s)
640
641
+ @extensible
642
async def unified_turn(
643
self,
644
system_message="",
tests/test_model_call_extensions.py
new
+34
@@ -0,0 +1,34 @@
1
+import pytest
2
+
3
+from helpers import extension
4
+from helpers.llm_result import LLMResult
5
+from models import LiteLLMChatWrapper
6
+
7
+
8
+@pytest.mark.asyncio
9
+@pytest.mark.parametrize(
10
+ ("method_name", "result"),
11
+ [
12
+ ("unified_call", ("response", "reasoning")),
13
+ ("unified_turn", LLMResult(response="response")),
14
+ ],
15
+)
16
+async def test_unified_model_calls_expose_function_extensions(
17
+ monkeypatch, method_name, result
18
+):
19
+ points = []
20
+
21
+ async def call_extensions(point, agent=None, **kwargs):
22
+ points.append(point)
23
+ if point.endswith("/start"):
24
+ kwargs["data"]["result"] = result
25
+
26
+ monkeypatch.setattr(extension, "call_extensions_async", call_extensions)
27
+
28
+ actual = await getattr(LiteLLMChatWrapper, method_name)(object())
29
+
30
+ assert actual is result
31
+ assert points == [
32
+ f"_functions/models/LiteLLMChatWrapper/{method_name}/start",
33
+ f"_functions/models/LiteLLMChatWrapper/{method_name}/end",
34
+ ]