Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings

Commit 1994f9d

Browse filesBrowse files
feat: pass extra_body through to LiteLLM acompletion (#638)
**Purpose** Allow arbitrary `extra_body` parameters (e.g. `cached_content`) to be forwarded into the LiteLLM call. Useful for context caching in Gemini models ([docs](https://ai.google.dev/gemini-api/docs/caching?lang=python)). **Example usage** ```python import os from agents import Agent, ModelSettings from agents.extensions.models.litellm_model import LitellmModel cache_name = "cachedContents/34jopukfx5di" # previously stored context gemini_model = LitellmModel( model="gemini/gemini-1.5-flash-002", api_key=os.getenv("GOOGLE_API_KEY") ) agent = Agent( name="Cached Gemini Agent", model=gemini_model, model_settings=ModelSettings( extra_body={"cached_content": cache_name} ) )
1 parent 2c46dae commit 1994f9d
Copy full SHA for 1994f9d

File tree

2 files changed

+47
-0
lines changed
Filter options

2 files changed

+47
-0
lines changed

‎src/agents/extensions/models/litellm_model.py

Copy file name to clipboardExpand all lines: src/agents/extensions/models/litellm_model.py
+2Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,8 @@ async def _fetch_response(
269269
extra_kwargs["extra_query"] = model_settings.extra_query
270270
if model_settings.metadata:
271271
extra_kwargs["metadata"] = model_settings.metadata
272+
if model_settings.extra_body and isinstance(model_settings.extra_body, dict):
273+
extra_kwargs.update(model_settings.extra_body)
272274

273275
ret = await litellm.acompletion(
274276
model=self.model,
+45Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
import litellm
2+
import pytest
3+
from litellm.types.utils import Choices, Message, ModelResponse, Usage
4+
5+
from agents.extensions.models.litellm_model import LitellmModel
6+
from agents.model_settings import ModelSettings
7+
from agents.models.interface import ModelTracing
8+
9+
10+
@pytest.mark.allow_call_model_methods
11+
@pytest.mark.asyncio
12+
async def test_extra_body_is_forwarded(monkeypatch):
13+
"""
14+
Forward `extra_body` entries into litellm.acompletion kwargs.
15+
16+
This ensures that user-provided parameters (e.g. cached_content)
17+
arrive alongside default arguments.
18+
"""
19+
captured: dict[str, object] = {}
20+
21+
async def fake_acompletion(model, messages=None, **kwargs):
22+
captured.update(kwargs)
23+
msg = Message(role="assistant", content="ok")
24+
choice = Choices(index=0, message=msg)
25+
return ModelResponse(choices=[choice], usage=Usage(0, 0, 0))
26+
27+
monkeypatch.setattr(litellm, "acompletion", fake_acompletion)
28+
settings = ModelSettings(
29+
temperature=0.1,
30+
extra_body={"cached_content": "some_cache", "foo": 123}
31+
)
32+
model = LitellmModel(model="test-model")
33+
34+
await model.get_response(
35+
system_instructions=None,
36+
input=[],
37+
model_settings=settings,
38+
tools=[],
39+
output_schema=None,
40+
handoffs=[],
41+
tracing=ModelTracing.DISABLED,
42+
previous_response_id=None,
43+
)
44+
45+
assert {"cached_content": "some_cache", "foo": 123}.items() <= captured.items()

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.