Merge pull request #4492 from Manouchehri/gemini-context-caching-1

feat(vertex_httpx.py): Support cachedContent.
This commit is contained in:
Krish Dholakia 2024-07-02 14:09:25 -07:00 committed by GitHub
commit 612af8f5be
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1035,6 +1035,9 @@ class VertexLLM(BaseLLM):
safety_settings: Optional[List[SafetSettingsConfig]] = optional_params.pop(
"safety_settings", None
) # type: ignore
cached_content: Optional[str] = optional_params.pop(
"cached_content", None
)
generation_config: Optional[GenerationConfig] = GenerationConfig(
**optional_params
)
@ -1050,6 +1053,8 @@ class VertexLLM(BaseLLM):
data["safetySettings"] = safety_settings
if generation_config is not None:
data["generationConfig"] = generation_config
if cached_content is not None:
data["cachedContent"] = cached_content
headers = {
"Content-Type": "application/json",