agno 2.3.8__py3-none-any.whl → 2.3.10__py3-none-any.whl

agno/media.py

@@ -4,6 +4,8 @@ from uuid import uuid4
 
 from pydantic import BaseModel, field_validator, model_validator
 
+from agno.utils.log import log_error
+
 
 class Image(BaseModel):
     """Unified Image class for all use cases (input, output, artifacts)"""
@@ -395,10 +397,20 @@ class File(BaseModel):
         name: Optional[str] = None,
         format: Optional[str] = None,
     ) -> "File":
-        """Create File from base64 encoded content"""
+        """Create File from base64 encoded content or plain text.
+
+        Handles both base64-encoded binary content and plain text content
+        (which is stored as UTF-8 strings for text/* MIME types).
+        """
         import base64
 
-        content_bytes = base64.b64decode(base64_content)
+        try:
+            content_bytes = base64.b64decode(base64_content)
+        except Exception:
+            # If not valid base64, it might be plain text content (text/csv, text/plain, etc.)
+            # which is stored as UTF-8 strings, not base64
+            content_bytes = base64_content.encode("utf-8")
+
         return cls(
             content=content_bytes,
             id=id,
@@ -413,10 +425,14 @@ class File(BaseModel):
         import httpx
 
         if self.url:
-            response = httpx.get(self.url)
-            content = response.content
-            mime_type = response.headers.get("Content-Type", "").split(";")[0]
-            return content, mime_type
+            try:
+                response = httpx.get(self.url)
+                content = response.content
+                mime_type = response.headers.get("Content-Type", "").split(";")[0]
+                return content, mime_type
+            except Exception:
+                log_error(f"Failed to download file from {self.url}")
+                return None
         else:
             return None
 

agno/models/aws/claude.py

@@ -7,7 +7,7 @@ from pydantic import BaseModel
 
 from agno.models.anthropic import Claude as AnthropicClaude
 from agno.utils.http import get_default_async_client, get_default_sync_client
-from agno.utils.log import log_debug, log_error, log_warning
+from agno.utils.log import log_debug, log_warning
 from agno.utils.models.claude import format_tools_for_model
 
 try:
@@ -70,8 +70,8 @@ class Claude(AnthropicClaude):
                 if self.aws_region:
                     client_params["aws_region"] = self.aws_region
             else:
-                self.aws_access_key = self.aws_access_key or getenv("AWS_ACCESS_KEY")
-                self.aws_secret_key = self.aws_secret_key or getenv("AWS_SECRET_KEY")
+                self.aws_access_key = self.aws_access_key or getenv("AWS_ACCESS_KEY_ID") or getenv("AWS_ACCESS_KEY")
+                self.aws_secret_key = self.aws_secret_key or getenv("AWS_SECRET_ACCESS_KEY") or getenv("AWS_SECRET_KEY")
                 self.aws_region = self.aws_region or getenv("AWS_REGION")
 
                 client_params = {
@@ -79,10 +79,11 @@ class Claude(AnthropicClaude):
                     "aws_access_key": self.aws_access_key,
                     "aws_region": self.aws_region,
                 }
-        if not (self.aws_access_key or (self.aws_access_key and self.aws_secret_key)):
-            log_error(
-                "AWS credentials not found. Please either set the AWS_BEDROCK_API_KEY or AWS_ACCESS_KEY and AWS_SECRET_KEY environment variables."
-            )
+
+            if not (self.api_key or (self.aws_access_key and self.aws_secret_key)):
+                log_warning(
+                    "AWS credentials not found. Please set AWS_BEDROCK_API_KEY or AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables or provide a boto3 session."
+                )
 
         if self.timeout is not None:
             client_params["timeout"] = self.timeout

agno/models/base.py

@@ -24,7 +24,7 @@ from uuid import uuid4
 
 from pydantic import BaseModel
 
-from agno.exceptions import AgentRunException, ModelProviderError
+from agno.exceptions import AgentRunException, ModelProviderError, RetryableModelProviderError
 from agno.media import Audio, File, Image, Video
 from agno.models.message import Citations, Message
 from agno.models.metrics import Metrics
@@ -153,6 +153,9 @@ class Model(ABC):
     delay_between_retries: int = 1
     # Exponential backoff: if True, the delay between retries is doubled each time
     exponential_backoff: bool = False
+    # Enable retrying a model invocation once with a guidance message.
+    # This is useful for known errors avoidable with extra instructions.
+    retry_with_guidance: bool = True
 
     def __post_init__(self):
         if self.provider is None and self.name is not None:
@@ -186,6 +189,9 @@ class Model(ABC):
                     sleep(delay)
                 else:
                     log_error(f"Model provider error after {self.retries + 1} attempts: {e}")
+            except RetryableModelProviderError as e:
+                kwargs["messages"].append(Message(role="user", content=e.retry_guidance_message, temporary=True))
+                return self._invoke_with_retry(**kwargs, retrying_with_guidance=True)
 
         # If we've exhausted all retries, raise the last exception
         raise last_exception  # type: ignore
@@ -212,6 +218,9 @@ class Model(ABC):
                     await asyncio.sleep(delay)
                 else:
                     log_error(f"Model provider error after {self.retries + 1} attempts: {e}")
+            except RetryableModelProviderError as e:
+                kwargs["messages"].append(Message(role="user", content=e.retry_guidance_message, temporary=True))
+                return await self._ainvoke_with_retry(**kwargs, retrying_with_guidance=True)
 
         # If we've exhausted all retries, raise the last exception
         raise last_exception  # type: ignore
@@ -240,6 +249,10 @@ class Model(ABC):
                     sleep(delay)
                 else:
                     log_error(f"Model provider error after {self.retries + 1} attempts: {e}")
+            except RetryableModelProviderError as e:
+                kwargs["messages"].append(Message(role="user", content=e.retry_guidance_message, temporary=True))
+                yield from self._invoke_stream_with_retry(**kwargs, retrying_with_guidance=True)
+                return  # Success, exit after regeneration
 
         # If we've exhausted all retries, raise the last exception
         raise last_exception  # type: ignore
@@ -269,6 +282,11 @@ class Model(ABC):
                     await asyncio.sleep(delay)
                 else:
                     log_error(f"Model provider error after {self.retries + 1} attempts: {e}")
+            except RetryableModelProviderError as e:
+                kwargs["messages"].append(Message(role="user", content=e.retry_guidance_message, temporary=True))
+                async for response in self._ainvoke_stream_with_retry(**kwargs, retrying_with_guidance=True):
+                    yield response
+                return  # Success, exit after regeneration
 
         # If we've exhausted all retries, raise the last exception
         raise last_exception  # type: ignore
@@ -278,6 +296,14 @@ class Model(ABC):
         _dict = {field: getattr(self, field) for field in fields if getattr(self, field) is not None}
         return _dict
 
+    def _remove_temporarys(self, messages: List[Message]) -> None:
+        """Remove temporal messages from the given list.
+
+        Args:
+            messages: The list of messages to filter (modified in place).
+        """
+        messages[:] = [m for m in messages if not m.temporary]
+
     def get_provider(self) -> str:
         return self.provider or self.name or self.__class__.__name__
 
@@ -1775,6 +1801,17 @@ class Model(ABC):
                 log_error(f"Error while iterating function result generator for {function_call.function.name}: {e}")
                 function_call.error = str(e)
                 function_call_success = False
+
+            # For generators, re-capture updated_session_state after consumption
+            # since session_state modifications were made during iteration
+            if function_execution_result.updated_session_state is None:
+                if (
+                    function_call.function._run_context is not None
+                    and function_call.function._run_context.session_state is not None
+                ):
+                    function_execution_result.updated_session_state = function_call.function._run_context.session_state
+                elif function_call.function._session_state is not None:
+                    function_execution_result.updated_session_state = function_call.function._session_state
         else:
             from agno.tools.function import ToolResult
 
@@ -2301,7 +2338,29 @@ class Model(ABC):
                     log_error(f"Error while iterating function result generator for {function_call.function.name}: {e}")
                     function_call.error = str(e)
                     function_call_success = False
-            else:
+
+            # For generators (sync or async), re-capture updated_session_state after consumption
+            # since session_state modifications were made during iteration
+            if async_function_call_output is not None or isinstance(
+                function_call.result,
+                (GeneratorType, collections.abc.Iterator, AsyncGeneratorType, collections.abc.AsyncIterator),
+            ):
+                if updated_session_state is None:
+                    if (
+                        function_call.function._run_context is not None
+                        and function_call.function._run_context.session_state is not None
+                    ):
+                        updated_session_state = function_call.function._run_context.session_state
+                    elif function_call.function._session_state is not None:
+                        updated_session_state = function_call.function._session_state
+
+            if not (
+                async_function_call_output is not None
+                or isinstance(
+                    function_call.result,
+                    (GeneratorType, collections.abc.Iterator, AsyncGeneratorType, collections.abc.AsyncIterator),
+                )
+            ):
                 from agno.tools.function import ToolResult
 
                 if isinstance(function_execution_result.result, ToolResult):

agno/models/deepseek/deepseek.py

@@ -3,7 +3,10 @@ from os import getenv
 from typing import Any, Dict, Optional
 
 from agno.exceptions import ModelAuthenticationError
+from agno.models.message import Message
 from agno.models.openai.like import OpenAILike
+from agno.utils.log import log_warning
+from agno.utils.openai import _format_file_for_message, audio_to_message, images_to_message
 
 
 @dataclass
@@ -58,3 +61,67 @@ class DeepSeek(OpenAILike):
         if self.client_params:
             client_params.update(self.client_params)
         return client_params
+
+    def _format_message(self, message: Message, compress_tool_results: bool = False) -> Dict[str, Any]:
+        """
+        Format a message into the format expected by OpenAI.
+
+        Args:
+            message (Message): The message to format.
+            compress_tool_results: Whether to compress tool results.
+
+        Returns:
+            Dict[str, Any]: The formatted message.
+        """
+        tool_result = message.get_content(use_compressed_content=compress_tool_results)
+
+        message_dict: Dict[str, Any] = {
+            "role": self.role_map[message.role] if self.role_map else self.default_role_map[message.role],
+            "content": tool_result,
+            "name": message.name,
+            "tool_call_id": message.tool_call_id,
+            "tool_calls": message.tool_calls,
+            "reasoning_content": message.reasoning_content,
+        }
+        message_dict = {k: v for k, v in message_dict.items() if v is not None}
+
+        # Ignore non-string message content
+        # because we assume that the images/audio are already added to the message
+        if (message.images is not None and len(message.images) > 0) or (
+            message.audio is not None and len(message.audio) > 0
+        ):
+            # Ignore non-string message content
+            # because we assume that the images/audio are already added to the message
+            if isinstance(message.content, str):
+                message_dict["content"] = [{"type": "text", "text": message.content}]
+                if message.images is not None:
+                    message_dict["content"].extend(images_to_message(images=message.images))
+
+                if message.audio is not None:
+                    message_dict["content"].extend(audio_to_message(audio=message.audio))
+
+        if message.audio_output is not None:
+            message_dict["content"] = ""
+            message_dict["audio"] = {"id": message.audio_output.id}
+
+        if message.videos is not None and len(message.videos) > 0:
+            log_warning("Video input is currently unsupported.")
+
+        if message.files is not None:
+            # Ensure content is a list of parts
+            content = message_dict.get("content")
+            if isinstance(content, str):  # wrap existing text
+                text = content
+                message_dict["content"] = [{"type": "text", "text": text}]
+            elif content is None:
+                message_dict["content"] = []
+            # Insert each file part before text parts
+            for file in message.files:
+                file_part = _format_file_for_message(file)
+                if file_part:
+                    message_dict["content"].insert(0, file_part)
+
+        # Manually add the content field even if it is None
+        if message.content is None:
+            message_dict["content"] = ""
+        return message_dict

agno/models/google/gemini.py

@@ -13,7 +13,8 @@ from pydantic import BaseModel
 
 from agno.exceptions import ModelProviderError
 from agno.media import Audio, File, Image, Video
-from agno.models.base import Model
+from agno.models.base import Model, RetryableModelProviderError
+from agno.models.google.utils import MALFORMED_FUNCTION_CALL_GUIDANCE, GeminiFinishReason
 from agno.models.message import Citations, Message, UrlCitation
 from agno.models.metrics import Metrics
 from agno.models.response import ModelResponse
@@ -35,6 +36,7 @@ try:
         GenerateContentResponseUsageMetadata,
         GoogleSearch,
         GoogleSearchRetrieval,
+        GroundingMetadata,
         Operation,
         Part,
         Retrieval,
@@ -243,8 +245,8 @@ class Gemini(Model):
         builtin_tools = []
 
         if self.grounding:
-            log_info(
-                "Grounding enabled. This is a legacy tool. For Gemini 2.0+ Please use enable `search` flag instead."
+            log_debug(
+                "Gemini Grounding enabled. This is a legacy tool. For Gemini 2.0+ Please use enable `search` flag instead."
             )
             builtin_tools.append(
                 Tool(
@@ -257,15 +259,15 @@ class Gemini(Model):
             )
 
         if self.search:
-            log_info("Google Search enabled.")
+            log_debug("Gemini Google Search enabled.")
             builtin_tools.append(Tool(google_search=GoogleSearch()))
 
         if self.url_context:
-            log_info("URL context enabled.")
+            log_debug("Gemini URL context enabled.")
             builtin_tools.append(Tool(url_context=UrlContext()))
 
         if self.vertexai_search:
-            log_info("Vertex AI Search enabled.")
+            log_debug("Gemini Vertex AI Search enabled.")
             if not self.vertexai_search_datastore:
                 log_error("vertexai_search_datastore must be provided when vertexai_search is enabled.")
                 raise ValueError("vertexai_search_datastore must be provided when vertexai_search is enabled.")
@@ -317,6 +319,7 @@ class Gemini(Model):
         tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
         run_response: Optional[RunOutput] = None,
         compress_tool_results: bool = False,
+        retrying_with_guidance: bool = False,
     ) -> ModelResponse:
         """
         Invokes the model with a list of messages and returns the response.
@@ -337,7 +340,13 @@ class Gemini(Model):
             )
             assistant_message.metrics.stop_timer()
 
-            model_response = self._parse_provider_response(provider_response, response_format=response_format)
+            model_response = self._parse_provider_response(
+                provider_response, response_format=response_format, retrying_with_guidance=retrying_with_guidance
+            )
+
+            # If we were retrying the invoke with guidance, remove the guidance message
+            if retrying_with_guidance is True:
+                self._remove_temporarys(messages)
 
             return model_response
 
@@ -350,6 +359,8 @@ class Gemini(Model):
                 model_name=self.name,
                 model_id=self.id,
             ) from e
+        except RetryableModelProviderError:
+            raise
         except Exception as e:
             log_error(f"Unknown error from Gemini API: {e}")
             raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
@@ -363,6 +374,7 @@ class Gemini(Model):
         tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
         run_response: Optional[RunOutput] = None,
         compress_tool_results: bool = False,
+        retrying_with_guidance: bool = False,
     ) -> Iterator[ModelResponse]:
         """
         Invokes the model with a list of messages and returns the response as a stream.
@@ -382,7 +394,11 @@ class Gemini(Model):
                 contents=formatted_messages,
                 **request_kwargs,
             ):
-                yield self._parse_provider_response_delta(response)
+                yield self._parse_provider_response_delta(response, retrying_with_guidance=retrying_with_guidance)
+
+            # If we were retrying the invoke with guidance, remove the guidance message
+            if retrying_with_guidance is True:
+                self._remove_temporarys(messages)
 
             assistant_message.metrics.stop_timer()
 
@@ -394,6 +410,8 @@ class Gemini(Model):
                 model_name=self.name,
                 model_id=self.id,
             ) from e
+        except RetryableModelProviderError:
+            raise
         except Exception as e:
             log_error(f"Unknown error from Gemini API: {e}")
             raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
@@ -407,6 +425,7 @@ class Gemini(Model):
         tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
         run_response: Optional[RunOutput] = None,
         compress_tool_results: bool = False,
+        retrying_with_guidance: bool = False,
     ) -> ModelResponse:
         """
         Invokes the model with a list of messages and returns the response.
@@ -429,7 +448,13 @@ class Gemini(Model):
             )
             assistant_message.metrics.stop_timer()
 
-            model_response = self._parse_provider_response(provider_response, response_format=response_format)
+            model_response = self._parse_provider_response(
+                provider_response, response_format=response_format, retrying_with_guidance=retrying_with_guidance
+            )
+
+            # If we were retrying the invoke with guidance, remove the guidance message
+            if retrying_with_guidance is True:
+                self._remove_temporarys(messages)
 
             return model_response
 
@@ -441,6 +466,8 @@ class Gemini(Model):
                 model_name=self.name,
                 model_id=self.id,
             ) from e
+        except RetryableModelProviderError:
+            raise
         except Exception as e:
             log_error(f"Unknown error from Gemini API: {e}")
             raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
@@ -454,6 +481,7 @@ class Gemini(Model):
         tool_choice: Optional[Union[str, Dict[str, Any]]] = None,
         run_response: Optional[RunOutput] = None,
         compress_tool_results: bool = False,
+        retrying_with_guidance: bool = False,
     ) -> AsyncIterator[ModelResponse]:
         """
         Invokes the model with a list of messages and returns the response as a stream.
@@ -476,7 +504,11 @@ class Gemini(Model):
                 **request_kwargs,
             )
             async for chunk in async_stream:
-                yield self._parse_provider_response_delta(chunk)
+                yield self._parse_provider_response_delta(chunk, retrying_with_guidance=retrying_with_guidance)
+
+            # If we were retrying the invoke with guidance, remove the guidance message
+            if retrying_with_guidance is True:
+                self._remove_temporarys(messages)
 
             assistant_message.metrics.stop_timer()
 
@@ -488,6 +520,8 @@ class Gemini(Model):
                 model_name=self.name,
                 model_id=self.id,
             ) from e
+        except RetryableModelProviderError:
+            raise
         except Exception as e:
             log_error(f"Unknown error from Gemini API: {e}")
             raise ModelProviderError(message=str(e), model_name=self.name, model_id=self.id) from e
@@ -863,10 +897,10 @@ class Gemini(Model):
 
     def _parse_provider_response(self, response: GenerateContentResponse, **kwargs) -> ModelResponse:
         """
-        Parse the OpenAI response into a ModelResponse.
+        Parse the Gemini response into a ModelResponse.
 
         Args:
-            response: Raw response from OpenAI
+            response: Raw response from Gemini
 
         Returns:
             ModelResponse: Parsed response data
@@ -875,8 +909,20 @@ class Gemini(Model):
 
         # Get response message
         response_message = Content(role="model", parts=[])
-        if response.candidates and response.candidates[0].content:
-            response_message = response.candidates[0].content
+        if response.candidates and len(response.candidates) > 0:
+            candidate = response.candidates[0]
+
+            # Raise if the request failed because of a malformed function call
+            if hasattr(candidate, "finish_reason") and candidate.finish_reason:
+                if candidate.finish_reason == GeminiFinishReason.MALFORMED_FUNCTION_CALL.value:
+                    # We only want to raise errors that trigger regeneration attempts once
+                    if kwargs.get("retrying_with_guidance") is True:
+                        pass
+                    if self.retry_with_guidance:
+                        raise RetryableModelProviderError(retry_guidance_message=MALFORMED_FUNCTION_CALL_GUIDANCE)
+
+            if candidate.content:
+                response_message = candidate.content
 
         # Add role
         if response_message.role is not None:
@@ -963,27 +1009,24 @@ class Gemini(Model):
             citations = Citations()
             citations_raw = {}
             citations_urls = []
+            web_search_queries: List[str] = []
 
             if response.candidates and response.candidates[0].grounding_metadata is not None:
-                grounding_metadata = response.candidates[0].grounding_metadata.model_dump()
-                citations_raw["grounding_metadata"] = grounding_metadata
+                grounding_metadata: GroundingMetadata = response.candidates[0].grounding_metadata
+                citations_raw["grounding_metadata"] = grounding_metadata.model_dump()
 
-                chunks = grounding_metadata.get("grounding_chunks", []) or []
-                citation_pairs = []
+                chunks = grounding_metadata.grounding_chunks or []
+                web_search_queries = grounding_metadata.web_search_queries or []
                 for chunk in chunks:
-                    if not isinstance(chunk, dict):
+                    if not chunk:
                         continue
-                    web = chunk.get("web")
-                    if not isinstance(web, dict):
+                    web = chunk.web
+                    if not web:
                         continue
-                    uri = web.get("uri")
-                    title = web.get("title")
+                    uri = web.uri
+                    title = web.title
                     if uri:
-                        citation_pairs.append((uri, title))
-
-                # Create citation objects from filtered pairs
-                grounding_urls = [UrlCitation(url=url, title=title) for url, title in citation_pairs]
-                citations_urls.extend(grounding_urls)
+                        citations_urls.append(UrlCitation(url=uri, title=title))
 
             # Handle URLs from URL context tool
             if (
@@ -991,22 +1034,29 @@ class Gemini(Model):
                 and hasattr(response.candidates[0], "url_context_metadata")
                 and response.candidates[0].url_context_metadata is not None
             ):
-                url_context_metadata = response.candidates[0].url_context_metadata.model_dump()
-                citations_raw["url_context_metadata"] = url_context_metadata
+                url_context_metadata = response.candidates[0].url_context_metadata
+                citations_raw["url_context_metadata"] = url_context_metadata.model_dump()
 
-                url_metadata_list = url_context_metadata.get("url_metadata", [])
+                url_metadata_list = url_context_metadata.url_metadata or []
                 for url_meta in url_metadata_list:
-                    retrieved_url = url_meta.get("retrieved_url")
-                    status = url_meta.get("url_retrieval_status", "UNKNOWN")
+                    retrieved_url = url_meta.retrieved_url
+                    status = "UNKNOWN"
+                    if url_meta.url_retrieval_status:
+                        status = url_meta.url_retrieval_status.value
                     if retrieved_url and status == "URL_RETRIEVAL_STATUS_SUCCESS":
                         # Avoid duplicate URLs
                         existing_urls = [citation.url for citation in citations_urls]
                         if retrieved_url not in existing_urls:
                             citations_urls.append(UrlCitation(url=retrieved_url, title=retrieved_url))
 
+            if citations_raw:
+                citations.raw = citations_raw
+            if citations_urls:
+                citations.urls = citations_urls
+            if web_search_queries:
+                citations.search_queries = web_search_queries
+
             if citations_raw or citations_urls:
-                citations.raw = citations_raw if citations_raw else None
-                citations.urls = citations_urls if citations_urls else None
                 model_response.citations = citations
 
         # Extract usage metadata if present
@@ -1019,11 +1069,20 @@ class Gemini(Model):
 
         return model_response
 
-    def _parse_provider_response_delta(self, response_delta: GenerateContentResponse) -> ModelResponse:
+    def _parse_provider_response_delta(self, response_delta: GenerateContentResponse, **kwargs) -> ModelResponse:
         model_response = ModelResponse()
 
         if response_delta.candidates and len(response_delta.candidates) > 0:
-            candidate_content = response_delta.candidates[0].content
+            candidate = response_delta.candidates[0]
+            candidate_content = candidate.content
+
+            # Raise if the request failed because of a malformed function call
+            if hasattr(candidate, "finish_reason") and candidate.finish_reason:
+                if candidate.finish_reason == GeminiFinishReason.MALFORMED_FUNCTION_CALL.value:
+                    if kwargs.get("retrying_with_guidance") is True:
+                        pass
+                    raise RetryableModelProviderError(retry_guidance_message=MALFORMED_FUNCTION_CALL_GUIDANCE)
+
             response_message: Content = Content(role="model", parts=[])
             if candidate_content is not None:
                 response_message = candidate_content
@@ -1096,28 +1155,52 @@ class Gemini(Model):
 
                         model_response.tool_calls.append(tool_call)
 
-            if response_delta.candidates[0].grounding_metadata is not None:
-                citations = Citations()
-                grounding_metadata = response_delta.candidates[0].grounding_metadata.model_dump()
-                citations.raw = grounding_metadata
+            citations = Citations()
+            citations.raw = {}
+            citations.urls = []
 
+            if (
+                hasattr(response_delta.candidates[0], "grounding_metadata")
+                and response_delta.candidates[0].grounding_metadata is not None
+            ):
+                grounding_metadata = response_delta.candidates[0].grounding_metadata
+                citations.raw["grounding_metadata"] = grounding_metadata.model_dump()
+                citations.search_queries = grounding_metadata.web_search_queries or []
                 # Extract url and title
-                chunks = grounding_metadata.pop("grounding_chunks", None) or []
-                citation_pairs = []
+                chunks = grounding_metadata.grounding_chunks or []
                 for chunk in chunks:
-                    if not isinstance(chunk, dict):
+                    if not chunk:
                         continue
-                    web = chunk.get("web")
-                    if not isinstance(web, dict):
+                    web = chunk.web
+                    if not web:
                         continue
-                    uri = web.get("uri")
-                    title = web.get("title")
+                    uri = web.uri
+                    title = web.title
                     if uri:
-                        citation_pairs.append((uri, title))
+                        citations.urls.append(UrlCitation(url=uri, title=title))
+
+            # Handle URLs from URL context tool
+            if (
+                hasattr(response_delta.candidates[0], "url_context_metadata")
+                and response_delta.candidates[0].url_context_metadata is not None
+            ):
+                url_context_metadata = response_delta.candidates[0].url_context_metadata
 
-                # Create citation objects from filtered pairs
-                citations.urls = [UrlCitation(url=url, title=title) for url, title in citation_pairs]
+                citations.raw["url_context_metadata"] = url_context_metadata.model_dump()
+
+                url_metadata_list = url_context_metadata.url_metadata or []
+                for url_meta in url_metadata_list:
+                    retrieved_url = url_meta.retrieved_url
+                    status = "UNKNOWN"
+                    if url_meta.url_retrieval_status:
+                        status = url_meta.url_retrieval_status.value
+                    if retrieved_url and status == "URL_RETRIEVAL_STATUS_SUCCESS":
+                        # Avoid duplicate URLs
+                        existing_urls = [citation.url for citation in citations.urls]
+                        if retrieved_url not in existing_urls:
+                            citations.urls.append(UrlCitation(url=retrieved_url, title=retrieved_url))
 
+            if citations.raw or citations.urls:
                 model_response.citations = citations
 
             # Extract usage metadata if present