google-adk 1.1.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

google/adk/models/lite_llm.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
 
 import base64
 import json
@@ -62,6 +63,7 @@ class FunctionChunk(BaseModel):
   id: Optional[str]
   name: Optional[str]
   args: Optional[str]
+  index: Optional[int] = 0
 
 
 class TextChunk(BaseModel):
@@ -136,7 +138,7 @@ def _safe_json_serialize(obj) -> str:
 
   try:
     # Try direct JSON serialization first
-    return json.dumps(obj)
+    return json.dumps(obj, ensure_ascii=False)
   except (TypeError, OverflowError):
     return str(obj)
 
@@ -186,7 +188,7 @@ def _content_to_message_param(
                 id=part.function_call.id,
                 function=Function(
                     name=part.function_call.name,
-                    arguments=json.dumps(part.function_call.args),
+                    arguments=_safe_json_serialize(part.function_call.args),
                 ),
             )
         )
@@ -194,6 +196,14 @@ def _content_to_message_param(
         content_present = True
 
     final_content = message_content if content_present else None
+    if final_content and isinstance(final_content, list):
+      # when the content is a single text object, we can use it directly.
+      # this is needed for ollama_chat provider which fails if content is a list
+      final_content = (
+          final_content[0].get("text", "")
+          if final_content[0].get("type", None) == "text"
+          else final_content
+      )
 
     return ChatCompletionAssistantMessage(
         role=role,
@@ -386,6 +396,7 @@ def _model_response_to_chunk(
               id=tool_call.id,
               name=tool_call.function.name,
               args=tool_call.function.arguments,
+              index=tool_call.index,
           ), finish_reason
 
     if finish_reason and not (
@@ -477,7 +488,7 @@ def _get_completion_inputs(
     llm_request: The LlmRequest to convert.
 
   Returns:
-    The litellm inputs (message list and tool dictionary).
+    The litellm inputs (message list, tool dictionary and response format).
   """
   messages = []
   for content in llm_request.contents or []:
@@ -506,7 +517,13 @@ def _get_completion_inputs(
         _function_declaration_to_tool_param(tool)
         for tool in llm_request.config.tools[0].function_declarations
     ]
-  return messages, tools
+
+  response_format = None
+
+  if llm_request.config.response_schema:
+    response_format = llm_request.config.response_schema
+
+  return messages, tools, response_format
 
 
 def _build_function_declaration_log(
@@ -643,33 +660,48 @@ class LiteLlm(BaseLlm):
     self._maybe_append_user_content(llm_request)
     logger.debug(_build_request_log(llm_request))
 
-    messages, tools = _get_completion_inputs(llm_request)
+    messages, tools, response_format = _get_completion_inputs(llm_request)
 
     completion_args = {
         "model": self.model,
         "messages": messages,
         "tools": tools,
+        "response_format": response_format,
     }
     completion_args.update(self._additional_args)
 
     if stream:
       text = ""
-      function_name = ""
-      function_args = ""
-      function_id = None
+      # Track function calls by index
+      function_calls = {}  # index -> {name, args, id}
       completion_args["stream"] = True
       aggregated_llm_response = None
       aggregated_llm_response_with_tool_call = None
       usage_metadata = None
-
+      fallback_index = 0
       for part in self.llm_client.completion(**completion_args):
         for chunk, finish_reason in _model_response_to_chunk(part):
           if isinstance(chunk, FunctionChunk):
+            index = chunk.index or fallback_index
+            if index not in function_calls:
+              function_calls[index] = {"name": "", "args": "", "id": None}
+
             if chunk.name:
-              function_name += chunk.name
+              function_calls[index]["name"] += chunk.name
             if chunk.args:
-              function_args += chunk.args
-            function_id = chunk.id or function_id
+              function_calls[index]["args"] += chunk.args
+
+              # check if args is completed (workaround for improper chunk
+              # indexing)
+              try:
+                json.loads(function_calls[index]["args"])
+                fallback_index += 1
+              except json.JSONDecodeError:
+                pass
+
+            function_calls[index]["id"] = (
+                chunk.id or function_calls[index]["id"] or str(index)
+            )
           elif isinstance(chunk, TextChunk):
             text += chunk.text
             yield _message_to_generate_content_response(
@@ -686,28 +718,33 @@ class LiteLlm(BaseLlm):
                 total_token_count=chunk.total_tokens,
             )
 
-          if finish_reason == "tool_calls" and function_id:
+          if (
+              finish_reason == "tool_calls" or finish_reason == "stop"
+          ) and function_calls:
+            tool_calls = []
+            for index, func_data in function_calls.items():
+              if func_data["id"]:
+                tool_calls.append(
+                    ChatCompletionMessageToolCall(
+                        type="function",
+                        id=func_data["id"],
+                        function=Function(
+                            name=func_data["name"],
+                            arguments=func_data["args"],
+                            index=index,
+                        ),
+                    )
+                )
             aggregated_llm_response_with_tool_call = (
                 _message_to_generate_content_response(
                     ChatCompletionAssistantMessage(
                         role="assistant",
                         content="",
-                        tool_calls=[
-                            ChatCompletionMessageToolCall(
-                                type="function",
-                                id=function_id,
-                                function=Function(
-                                    name=function_name,
-                                    arguments=function_args,
-                                ),
-                            )
-                        ],
+                        tool_calls=tool_calls,
                     )
                 )
             )
-            function_name = ""
-            function_args = ""
-            function_id = None
+            function_calls.clear()
           elif finish_reason == "stop" and text:
             aggregated_llm_response = _message_to_generate_content_response(
                 ChatCompletionAssistantMessage(role="assistant", content=text)

google/adk/sessions/_session_util.py

@@ -11,34 +11,28 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 """Utility functions for session service."""
+from __future__ import annotations
 
-import base64
 from typing import Any
 from typing import Optional
 
 from google.genai import types
 
 
-def encode_content(content: types.Content):
-  """Encodes a content object to a JSON dictionary."""
-  encoded_content = content.model_dump(exclude_none=True)
-  for p in encoded_content["parts"]:
-    if "inline_data" in p:
-      p["inline_data"]["data"] = base64.b64encode(
-          p["inline_data"]["data"]
-      ).decode("utf-8")
-  return encoded_content
-
-
 def decode_content(
     content: Optional[dict[str, Any]],
 ) -> Optional[types.Content]:
   """Decodes a content object from a JSON dictionary."""
   if not content:
     return None
-  for p in content["parts"]:
-    if "inline_data" in p:
-      p["inline_data"]["data"] = base64.b64decode(p["inline_data"]["data"])
   return types.Content.model_validate(content)
+
+
+def decode_grounding_metadata(
+    grounding_metadata: Optional[dict[str, Any]],
+) -> Optional[types.GroundingMetadata]:
+  """Decodes a grounding metadata object from a JSON dictionary."""
+  if not grounding_metadata:
+    return None
+  return types.GroundingMetadata.model_validate(grounding_metadata)

google/adk/sessions/database_session_service.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
 import copy
 from datetime import datetime
 import json
@@ -19,6 +21,7 @@ from typing import Any
 from typing import Optional
 import uuid
 
+from google.genai import types
 from sqlalchemy import Boolean
 from sqlalchemy import delete
 from sqlalchemy import Dialect
@@ -89,6 +92,18 @@ class DynamicJSON(TypeDecorator):
     return value
 
 
+class PreciseTimestamp(TypeDecorator):
+  """Represents a timestamp precise to the microsecond."""
+
+  impl = DateTime
+  cache_ok = True
+
+  def load_dialect_impl(self, dialect):
+    if dialect.name == "mysql":
+      return dialect.type_descriptor(mysql.DATETIME(fsp=6))
+    return self.impl
+
+
 class Base(DeclarativeBase):
   """Base class for database tables."""
 
@@ -153,7 +168,9 @@ class StorageEvent(Base):
   branch: Mapped[str] = mapped_column(
       String(DEFAULT_MAX_VARCHAR_LENGTH), nullable=True
   )
-  timestamp: Mapped[DateTime] = mapped_column(DateTime(), default=func.now())
+  timestamp: Mapped[PreciseTimestamp] = mapped_column(
+      PreciseTimestamp, default=func.now()
+  )
   content: Mapped[dict[str, Any]] = mapped_column(DynamicJSON, nullable=True)
   actions: Mapped[MutableDict[str, Any]] = mapped_column(PickleType)
 
@@ -199,6 +216,55 @@ class StorageEvent(Base):
     else:
       self.long_running_tool_ids_json = json.dumps(list(value))
 
+  @classmethod
+  def from_event(cls, session: Session, event: Event) -> StorageEvent:
+    storage_event = StorageEvent(
+        id=event.id,
+        invocation_id=event.invocation_id,
+        author=event.author,
+        branch=event.branch,
+        actions=event.actions,
+        session_id=session.id,
+        app_name=session.app_name,
+        user_id=session.user_id,
+        timestamp=datetime.fromtimestamp(event.timestamp),
+        long_running_tool_ids=event.long_running_tool_ids,
+        partial=event.partial,
+        turn_complete=event.turn_complete,
+        error_code=event.error_code,
+        error_message=event.error_message,
+        interrupted=event.interrupted,
+    )
+    if event.content:
+      storage_event.content = event.content.model_dump(
+          exclude_none=True, mode="json"
+      )
+    if event.grounding_metadata:
+      storage_event.grounding_metadata = event.grounding_metadata.model_dump(
+          exclude_none=True, mode="json"
+      )
+    return storage_event
+
+  def to_event(self) -> Event:
+    return Event(
+        id=self.id,
+        invocation_id=self.invocation_id,
+        author=self.author,
+        branch=self.branch,
+        actions=self.actions,
+        timestamp=self.timestamp.timestamp(),
+        content=_session_util.decode_content(self.content),
+        long_running_tool_ids=self.long_running_tool_ids,
+        partial=self.partial,
+        turn_complete=self.turn_complete,
+        error_code=self.error_code,
+        error_message=self.error_message,
+        interrupted=self.interrupted,
+        grounding_metadata=_session_util.decode_grounding_metadata(
+            self.grounding_metadata
+        ),
+    )
+
 
 class StorageAppState(Base):
   """Represents an app state stored in the database."""
@@ -238,14 +304,14 @@ class StorageUserState(Base):
 class DatabaseSessionService(BaseSessionService):
   """A session service that uses a database for storage."""
 
-  def __init__(self, db_url: str):
+  def __init__(self, db_url: str, **kwargs: Any):
     """Initializes the database session service with a database URL."""
     # 1. Create DB engine for db connection
     # 2. Create all tables based on schema
     # 3. Initialize all properties
 
     try:
-      db_engine = create_engine(db_url)
+      db_engine = create_engine(db_url, **kwargs)
     except Exception as e:
       if isinstance(e, ArgumentError):
         raise ValueError(
@@ -409,25 +475,7 @@ class DatabaseSessionService(BaseSessionService):
           state=merged_state,
           last_update_time=storage_session.update_time.timestamp(),
       )
-      session.events = [
-          Event(
-              id=e.id,
-              author=e.author,
-              branch=e.branch,
-              invocation_id=e.invocation_id,
-              content=_session_util.decode_content(e.content),
-              actions=e.actions,
-              timestamp=e.timestamp.timestamp(),
-              long_running_tool_ids=e.long_running_tool_ids,
-              grounding_metadata=e.grounding_metadata,
-              partial=e.partial,
-              turn_complete=e.turn_complete,
-              error_code=e.error_code,
-              error_message=e.error_message,
-              interrupted=e.interrupted,
-          )
-          for e in reversed(storage_events)
-      ]
+      session.events = [e.to_event() for e in reversed(storage_events)]
     return session
 
   @override
@@ -512,38 +560,18 @@ class DatabaseSessionService(BaseSessionService):
               _extract_state_delta(event.actions.state_delta)
           )
 
-      # Merge state
-      app_state.update(app_state_delta)
-      user_state.update(user_state_delta)
-      session_state.update(session_state_delta)
-
-      # Update storage
-      storage_app_state.state = app_state
-      storage_user_state.state = user_state
-      storage_session.state = session_state
-
-      storage_event = StorageEvent(
-          id=event.id,
-          invocation_id=event.invocation_id,
-          author=event.author,
-          branch=event.branch,
-          actions=event.actions,
-          session_id=session.id,
-          app_name=session.app_name,
-          user_id=session.user_id,
-          timestamp=datetime.fromtimestamp(event.timestamp),
-          long_running_tool_ids=event.long_running_tool_ids,
-          grounding_metadata=event.grounding_metadata,
-          partial=event.partial,
-          turn_complete=event.turn_complete,
-          error_code=event.error_code,
-          error_message=event.error_message,
-          interrupted=event.interrupted,
-      )
-      if event.content:
-        storage_event.content = _session_util.encode_content(event.content)
+      # Merge state and update storage
+      if app_state_delta:
+        app_state.update(app_state_delta)
+        storage_app_state.state = app_state
+      if user_state_delta:
+        user_state.update(user_state_delta)
+        storage_user_state.state = user_state
+      if session_state_delta:
+        session_state.update(session_state_delta)
+        storage_session.state = session_state
 
-      session_factory.add(storage_event)
+      session_factory.add(StorageEvent.from_event(session, event))
 
       session_factory.commit()
       session_factory.refresh(storage_session)
@@ -556,19 +584,6 @@ class DatabaseSessionService(BaseSessionService):
     return event
 
 
-def convert_event(event: StorageEvent) -> Event:
-  """Converts a storage event to an event."""
-  return Event(
-      id=event.id,
-      author=event.author,
-      branch=event.branch,
-      invocation_id=event.invocation_id,
-      content=event.content,
-      actions=event.actions,
-      timestamp=event.timestamp.timestamp(),
-  )
-
-
 def _extract_state_delta(state: dict[str, Any]):
   app_state_delta = {}
   user_state_delta = {}

google/adk/sessions/vertex_ai_session_service.py

@@ -11,13 +11,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
 import asyncio
 import logging
 import re
+import time
 from typing import Any
 from typing import Optional
+import urllib.parse
 
 from dateutil import parser
+from google.genai import types
 from typing_extensions import override
 
 from google import genai
@@ -154,15 +159,29 @@ class VertexAiSessionService(BaseSessionService):
     if list_events_api_response.get('httpHeaders', None):
       return session
 
-    session.events = [
+    session.events += [
         _from_api_event(event)
         for event in list_events_api_response['sessionEvents']
     ]
+
+    while list_events_api_response.get('nextPageToken', None):
+      page_token = list_events_api_response.get('nextPageToken', None)
+      list_events_api_response = await api_client.async_request(
+          http_method='GET',
+          path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}/events?pageToken={page_token}',
+          request_dict={},
+      )
+      session.events += [
+          _from_api_event(event)
+          for event in list_events_api_response['sessionEvents']
+      ]
+
     session.events = [
         event for event in session.events if event.timestamp <= update_timestamp
     ]
     session.events.sort(key=lambda event: event.timestamp)
 
+    # Filter events based on config
     if config:
       if config.num_recent_events:
         session.events = session.events[-config.num_recent_events :]
@@ -183,10 +202,15 @@ class VertexAiSessionService(BaseSessionService):
   ) -> ListSessionsResponse:
     reasoning_engine_id = _parse_reasoning_engine_id(app_name)
 
+    path = f'reasoningEngines/{reasoning_engine_id}/sessions'
+    if user_id:
+      parsed_user_id = urllib.parse.quote(f'''"{user_id}"''', safe='')
+      path = path + f'?filter=user_id={parsed_user_id}'
+
     api_client = _get_api_client(self.project, self.location)
     api_response = await api_client.async_request(
         http_method='GET',
-        path=f'reasoningEngines/{reasoning_engine_id}/sessions?filter=user_id={user_id}',
+        path=path,
         request_dict={},
     )
 
@@ -256,7 +280,7 @@ def _convert_event_to_json(event: Event):
   }
   if event.grounding_metadata:
     metadata_json['grounding_metadata'] = event.grounding_metadata.model_dump(
-        exclude_none=True
+        exclude_none=True, mode='json'
     )
 
   event_json = {
@@ -284,7 +308,9 @@ def _convert_event_to_json(event: Event):
     }
     event_json['actions'] = actions_json
   if event.content:
-    event_json['content'] = _session_util.encode_content(event.content)
+    event_json['content'] = event.content.model_dump(
+        exclude_none=True, mode='json'
+    )
   if event.error_code:
     event_json['error_code'] = event.error_code
   if event.error_message:
@@ -325,8 +351,8 @@ def _from_api_event(api_event: dict) -> Event:
     event.turn_complete = api_event['eventMetadata'].get('turnComplete', None)
     event.interrupted = api_event['eventMetadata'].get('interrupted', None)
     event.branch = api_event['eventMetadata'].get('branch', None)
-    event.grounding_metadata = api_event['eventMetadata'].get(
-        'groundingMetadata', None
+    event.grounding_metadata = _session_util.decode_grounding_metadata(
+        api_event['eventMetadata'].get('groundingMetadata', None)
     )
     event.long_running_tool_ids = (
         set(long_running_tool_ids_list) if long_running_tool_ids_list else None