From 25aeb1fa0019adc5a6a14a7d907ba7d5aa328d2b Mon Sep 17 00:00:00 2001 From: Nathan DeMoss Date: Wed, 26 Nov 2025 14:30:23 -0500 Subject: [PATCH] Enhance OpenAI response model with context size and location Added user_location and search_context_size fields with validation. Implemented context_window adjustment based on search_context_size. --- src/llama_stack_api/openai_responses.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/src/llama_stack_api/openai_responses.py b/src/llama_stack_api/openai_responses.py index e20004487..e4aebfab0 100644 --- a/src/llama_stack_api/openai_responses.py +++ b/src/llama_stack_api/openai_responses.py @@ -416,14 +416,28 @@ class OpenAIResponseInputToolWebSearch(BaseModel): # Must match values of WebSearchToolTypes above type: ( - Literal["web_search"] + Literal["web_search"] | Literal["web_search_preview"] | Literal["web_search_preview_2025_03_11"] | Literal["web_search_2025_08_26"] ) = "web_search" - # TODO: actually use search_context_size somewhere... + # Context size hint; validated to accepted values. Used to set an internal context window. search_context_size: str | None = Field(default="medium", pattern="^low|medium|high$") - # TODO: add user_location + # Optional user location hint to improve local relevance (e.g., "Seattle, WA"). + user_location: str | None = None + # Derived internal-only: token budget influenced by search_context_size. Excluded from wire schema. + context_window: int = Field(default=1024, exclude=True) + + @model_validator(mode="after") + def _apply_context_size(self) -> "OpenAIResponseInputToolWebSearch": + size = (self.search_context_size or "medium").lower() + if size == "low": + object.__setattr__(self, "context_window", 512) + elif size == "high": + object.__setattr__(self, "context_window", 2048) + else: + object.__setattr__(self, "context_window", 1024) + return self @json_schema_type