agno 2.2.5__py3-none-any.whl → 2.2.6__py3-none-any.whl

agno/os/schema.py

@@ -13,14 +13,17 @@ from agno.os.utils import (
     extract_input_media,
     format_team_tools,
     format_tools,
+    get_agent_input_schema_dict,
     get_run_input,
     get_session_name,
+    get_team_input_schema_dict,
     get_workflow_input_schema_dict,
 )
 from agno.run.agent import RunOutput
 from agno.run.team import TeamRunOutput
 from agno.session import AgentSession, TeamSession, WorkflowSession
 from agno.team.team import Team
+from agno.workflow.agent import WorkflowAgent
 from agno.workflow.workflow import Workflow
 
 
@@ -166,21 +169,25 @@ class ModelResponse(BaseModel):
 
 
 class AgentResponse(BaseModel):
-    id: Optional[str] = Field(None, description="Unique identifier for the agent")
-    name: Optional[str] = Field(None, description="Name of the agent")
-    db_id: Optional[str] = Field(None, description="Database identifier")
-    model: Optional[ModelResponse] = Field(None, description="Model configuration")
-    tools: Optional[Dict[str, Any]] = Field(None, description="Tool configurations")
-    sessions: Optional[Dict[str, Any]] = Field(None, description="Session configurations")
-    knowledge: Optional[Dict[str, Any]] = Field(None, description="Knowledge base configurations")
-    memory: Optional[Dict[str, Any]] = Field(None, description="Memory configurations")
-    reasoning: Optional[Dict[str, Any]] = Field(None, description="Reasoning configurations")
-    default_tools: Optional[Dict[str, Any]] = Field(None, description="Default tool settings")
-    system_message: Optional[Dict[str, Any]] = Field(None, description="System message configurations")
-    extra_messages: Optional[Dict[str, Any]] = Field(None, description="Extra message configurations")
-    response_settings: Optional[Dict[str, Any]] = Field(None, description="Response settings")
-    streaming: Optional[Dict[str, Any]] = Field(None, description="Streaming configurations")
-    metadata: Optional[Dict[str, Any]] = Field(None, description="Additional metadata")
+    id: Optional[str] = None
+    name: Optional[str] = None
+    db_id: Optional[str] = None
+    model: Optional[ModelResponse] = None
+    tools: Optional[Dict[str, Any]] = None
+    sessions: Optional[Dict[str, Any]] = None
+    knowledge: Optional[Dict[str, Any]] = None
+    memory: Optional[Dict[str, Any]] = None
+    reasoning: Optional[Dict[str, Any]] = None
+    default_tools: Optional[Dict[str, Any]] = None
+    system_message: Optional[Dict[str, Any]] = None
+    extra_messages: Optional[Dict[str, Any]] = None
+    response_settings: Optional[Dict[str, Any]] = None
+    streaming: Optional[Dict[str, Any]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    input_schema: Optional[Dict[str, Any]] = None
+
+    class Config:
+        exclude_none = True
 
     @classmethod
     async def from_agent(cls, agent: Agent) -> "AgentResponse":
@@ -375,6 +382,7 @@ class AgentResponse(BaseModel):
             "stream_events": agent.stream_events,
             "stream_intermediate_steps": agent.stream_intermediate_steps,
         }
+
         return AgentResponse(
             id=agent.id,
             name=agent.name,
@@ -391,28 +399,28 @@ class AgentResponse(BaseModel):
             response_settings=filter_meaningful_config(response_settings_info, agent_defaults),
             streaming=filter_meaningful_config(streaming_info, agent_defaults),
             metadata=agent.metadata,
+            input_schema=get_agent_input_schema_dict(agent),
         )
 
 
 class TeamResponse(BaseModel):
-    id: Optional[str] = Field(None, description="Unique identifier for the team")
-    name: Optional[str] = Field(None, description="Name of the team")
-    db_id: Optional[str] = Field(None, description="Database identifier")
-    description: Optional[str] = Field(None, description="Description of the team")
-    model: Optional[ModelResponse] = Field(None, description="Model configuration")
-    tools: Optional[Dict[str, Any]] = Field(None, description="Tool configurations")
-    sessions: Optional[Dict[str, Any]] = Field(None, description="Session configurations")
-    knowledge: Optional[Dict[str, Any]] = Field(None, description="Knowledge base configurations")
-    memory: Optional[Dict[str, Any]] = Field(None, description="Memory configurations")
-    reasoning: Optional[Dict[str, Any]] = Field(None, description="Reasoning configurations")
-    default_tools: Optional[Dict[str, Any]] = Field(None, description="Default tool settings")
-    system_message: Optional[Dict[str, Any]] = Field(None, description="System message configurations")
-    response_settings: Optional[Dict[str, Any]] = Field(None, description="Response settings")
-    streaming: Optional[Dict[str, Any]] = Field(None, description="Streaming configurations")
-    members: Optional[List[Union[AgentResponse, "TeamResponse"]]] = Field(
-        None, description="List of team members (agents or teams)"
-    )
-    metadata: Optional[Dict[str, Any]] = Field(None, description="Additional metadata")
+    id: Optional[str] = None
+    name: Optional[str] = None
+    db_id: Optional[str] = None
+    description: Optional[str] = None
+    model: Optional[ModelResponse] = None
+    tools: Optional[Dict[str, Any]] = None
+    sessions: Optional[Dict[str, Any]] = None
+    knowledge: Optional[Dict[str, Any]] = None
+    memory: Optional[Dict[str, Any]] = None
+    reasoning: Optional[Dict[str, Any]] = None
+    default_tools: Optional[Dict[str, Any]] = None
+    system_message: Optional[Dict[str, Any]] = None
+    response_settings: Optional[Dict[str, Any]] = None
+    streaming: Optional[Dict[str, Any]] = None
+    members: Optional[List[Union[AgentResponse, "TeamResponse"]]] = None
+    metadata: Optional[Dict[str, Any]] = None
+    input_schema: Optional[Dict[str, Any]] = None
 
     @classmethod
     async def from_team(cls, team: Team) -> "TeamResponse":
@@ -621,6 +629,7 @@ class TeamResponse(BaseModel):
             streaming=filter_meaningful_config(streaming_info, team_defaults),
             members=members if members else None,
             metadata=team.metadata,
+            input_schema=get_team_input_schema_dict(team),
         )
 
 
@@ -634,6 +643,7 @@ class WorkflowResponse(BaseModel):
     agent: Optional[AgentResponse] = Field(None, description="Agent configuration if used")
     team: Optional[TeamResponse] = Field(None, description="Team configuration if used")
     metadata: Optional[Dict[str, Any]] = Field(None, description="Additional metadata")
+    workflow_agent: bool = Field(False, description="Whether this workflow uses a WorkflowAgent")
 
     @classmethod
     async def _resolve_agents_and_teams_recursively(cls, steps: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
@@ -685,6 +695,7 @@ class WorkflowResponse(BaseModel):
             steps=steps,
             input_schema=get_workflow_input_schema_dict(workflow),
             metadata=workflow.metadata,
+            workflow_agent=isinstance(workflow.agent, WorkflowAgent) if workflow.agent else False,
         )
 
 
@@ -870,6 +881,7 @@ class RunSchema(BaseModel):
     created_at: Optional[datetime] = Field(None, description="Run creation timestamp")
     references: Optional[List[dict]] = Field(None, description="References cited in the run")
     reasoning_messages: Optional[List[dict]] = Field(None, description="Reasoning process messages")
+    session_state: Optional[dict] = Field(None, description="Session state at the end of the run")
     images: Optional[List[dict]] = Field(None, description="Images included in the run")
     videos: Optional[List[dict]] = Field(None, description="Videos included in the run")
     audio: Optional[List[dict]] = Field(None, description="Audio files included in the run")
@@ -897,6 +909,7 @@ class RunSchema(BaseModel):
             events=[event for event in run_dict["events"]] if run_dict.get("events") else None,
             references=run_dict.get("references", []),
             reasoning_messages=run_dict.get("reasoning_messages", []),
+            session_state=run_dict.get("session_state"),
             images=run_dict.get("images", []),
             videos=run_dict.get("videos", []),
             audio=run_dict.get("audio", []),
@@ -925,6 +938,7 @@ class TeamRunSchema(BaseModel):
     created_at: Optional[datetime] = Field(None, description="Run creation timestamp")
     references: Optional[List[dict]] = Field(None, description="References cited in the run")
     reasoning_messages: Optional[List[dict]] = Field(None, description="Reasoning process messages")
+    session_state: Optional[dict] = Field(None, description="Session state at the end of the run")
     input_media: Optional[Dict[str, Any]] = Field(None, description="Input media attachments")
     images: Optional[List[dict]] = Field(None, description="Images included in the run")
     videos: Optional[List[dict]] = Field(None, description="Videos included in the run")
@@ -954,6 +968,7 @@ class TeamRunSchema(BaseModel):
             else None,
             references=run_dict.get("references", []),
             reasoning_messages=run_dict.get("reasoning_messages", []),
+            session_state=run_dict.get("session_state"),
             images=run_dict.get("images", []),
             videos=run_dict.get("videos", []),
             audio=run_dict.get("audio", []),

agno/os/utils.py

@@ -274,6 +274,33 @@ def get_workflow_by_id(workflow_id: str, workflows: Optional[List[Workflow]] = N
     return None
 
 
+#  INPUT SCHEMA VALIDATIONS
+
+
+def get_agent_input_schema_dict(agent: Agent) -> Optional[Dict[str, Any]]:
+    """Get input schema as dictionary for API responses"""
+
+    if agent.input_schema is not None:
+        try:
+            return agent.input_schema.model_json_schema()
+        except Exception:
+            return None
+
+    return None
+
+
+def get_team_input_schema_dict(team: Team) -> Optional[Dict[str, Any]]:
+    """Get input schema as dictionary for API responses"""
+
+    if team.input_schema is not None:
+        try:
+            return team.input_schema.model_json_schema()
+        except Exception:
+            return None
+
+    return None
+
+
 def get_workflow_input_schema_dict(workflow: Workflow) -> Optional[Dict[str, Any]]:
     """Get input schema as dictionary for API responses"""
 

agno/run/agent.py

@@ -221,6 +221,9 @@ class RunContentEvent(BaseAgentRunEvent):
 
     event: str = RunEvent.run_content.value
     content: Optional[Any] = None
+    workflow_agent: bool = (
+        False  # Used by consumers of the events to distinguish between workflow agent and regular agent
+    )
     content_type: str = "str"
     reasoning_content: Optional[str] = None
     model_provider_data: Optional[Dict[str, Any]] = None
@@ -263,6 +266,7 @@ class RunCompletedEvent(BaseAgentRunEvent):
     reasoning_messages: Optional[List[Message]] = None
     metadata: Optional[Dict[str, Any]] = None
     metrics: Optional[Metrics] = None
+    session_state: Optional[Dict[str, Any]] = None
 
 
 @dataclass
@@ -527,6 +531,7 @@ class RunOutput:
     references: Optional[List[MessageReferences]] = None
 
     metadata: Optional[Dict[str, Any]] = None
+    session_state: Optional[Dict[str, Any]] = None
 
     created_at: int = field(default_factory=lambda: int(time()))
 

agno/run/team.py

@@ -261,6 +261,7 @@ class RunCompletedEvent(BaseTeamRunEvent):
     member_responses: List[Union["TeamRunOutput", RunOutput]] = field(default_factory=list)
     metadata: Optional[Dict[str, Any]] = None
     metrics: Optional[Metrics] = None
+    session_state: Optional[Dict[str, Any]] = None
 
 
 @dataclass
@@ -500,6 +501,7 @@ class TeamRunOutput:
     citations: Optional[Citations] = None
     model_provider_data: Optional[Dict[str, Any]] = None
     metadata: Optional[Dict[str, Any]] = None
+    session_state: Optional[Dict[str, Any]] = None
 
     references: Optional[List[MessageReferences]] = None
     additional_input: Optional[List[Message]] = None

agno/run/workflow.py

@@ -31,6 +31,9 @@ class WorkflowRunEvent(str, Enum):
     workflow_cancelled = "WorkflowCancelled"
     workflow_error = "WorkflowError"
 
+    workflow_agent_started = "WorkflowAgentStarted"
+    workflow_agent_completed = "WorkflowAgentCompleted"
+
     step_started = "StepStarted"
     step_completed = "StepCompleted"
     step_error = "StepError"
@@ -126,6 +129,21 @@ class WorkflowStartedEvent(BaseWorkflowRunOutputEvent):
     event: str = WorkflowRunEvent.workflow_started.value
 
 
+@dataclass
+class WorkflowAgentStartedEvent(BaseWorkflowRunOutputEvent):
+    """Event sent when workflow agent starts (before deciding to run workflow or answer directly)"""
+
+    event: str = WorkflowRunEvent.workflow_agent_started.value
+
+
+@dataclass
+class WorkflowAgentCompletedEvent(BaseWorkflowRunOutputEvent):
+    """Event sent when workflow agent completes (after running workflow or answering directly)"""
+
+    event: str = WorkflowRunEvent.workflow_agent_completed.value
+    content: Optional[Any] = None
+
+
 @dataclass
 class WorkflowCompletedEvent(BaseWorkflowRunOutputEvent):
     """Event sent when workflow execution completes"""
@@ -403,6 +421,8 @@ class CustomEvent(BaseWorkflowRunOutputEvent):
 # Union type for all workflow run response events
 WorkflowRunOutputEvent = Union[
     WorkflowStartedEvent,
+    WorkflowAgentStartedEvent,
+    WorkflowAgentCompletedEvent,
     WorkflowCompletedEvent,
     WorkflowErrorEvent,
     WorkflowCancelledEvent,
@@ -428,6 +448,8 @@ WorkflowRunOutputEvent = Union[
 # Map event string to dataclass for workflow events
 WORKFLOW_RUN_EVENT_TYPE_REGISTRY = {
     WorkflowRunEvent.workflow_started.value: WorkflowStartedEvent,
+    WorkflowRunEvent.workflow_agent_started.value: WorkflowAgentStartedEvent,
+    WorkflowRunEvent.workflow_agent_completed.value: WorkflowAgentCompletedEvent,
     WorkflowRunEvent.workflow_completed.value: WorkflowCompletedEvent,
     WorkflowRunEvent.workflow_cancelled.value: WorkflowCancelledEvent,
     WorkflowRunEvent.workflow_error.value: WorkflowErrorEvent,
@@ -491,6 +513,10 @@ class WorkflowRunOutput:
     # Store agent/team responses separately with parent_run_id references
     step_executor_runs: Optional[List[Union[RunOutput, TeamRunOutput]]] = None
 
+    # Workflow agent run - stores the full agent RunOutput when workflow agent is used
+    # The agent's parent_run_id will point to this workflow run's run_id to establish the relationship
+    workflow_agent_run: Optional[RunOutput] = None
+
     # Store events from workflow execution
     events: Optional[List[WorkflowRunOutputEvent]] = None
 
@@ -522,6 +548,7 @@ class WorkflowRunOutput:
                 "step_executor_runs",
                 "events",
                 "metrics",
+                "workflow_agent_run",
             ]
         }
 
@@ -557,6 +584,9 @@ class WorkflowRunOutput:
         if self.step_executor_runs:
             _dict["step_executor_runs"] = [run.to_dict() for run in self.step_executor_runs]
 
+        if self.workflow_agent_run is not None:
+            _dict["workflow_agent_run"] = self.workflow_agent_run.to_dict()
+
         if self.metrics is not None:
             _dict["metrics"] = self.metrics.to_dict()
 
@@ -604,6 +634,14 @@ class WorkflowRunOutput:
                 else:
                     step_executor_runs.append(RunOutput.from_dict(run_data))
 
+        workflow_agent_run_data = data.pop("workflow_agent_run", None)
+        workflow_agent_run = None
+        if workflow_agent_run_data:
+            if isinstance(workflow_agent_run_data, dict):
+                workflow_agent_run = RunOutput.from_dict(workflow_agent_run_data)
+            elif isinstance(workflow_agent_run_data, RunOutput):
+                workflow_agent_run = workflow_agent_run_data
+
         metadata = data.pop("metadata", None)
 
         images = reconstruct_images(data.pop("images", []))
@@ -640,6 +678,7 @@ class WorkflowRunOutput:
 
         return cls(
             step_results=parsed_step_results,
+            workflow_agent_run=workflow_agent_run,
             metadata=metadata,
             images=images,
             videos=videos,

agno/session/summary.py

@@ -1,11 +1,12 @@
 from dataclasses import dataclass
 from datetime import datetime
 from textwrap import dedent
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union
 
 from pydantic import BaseModel, Field
 
 from agno.models.base import Model
+from agno.models.utils import get_model
 from agno.run.agent import Message
 from agno.utils.log import log_debug, log_warning
 
@@ -142,7 +143,10 @@ class SessionSummaryManager:
         if not session:
             return None
 
-        self.model = cast(Model, self.model)
+        self.model = get_model(self.model)
+        if self.model is None:
+            return None
+
         response_format = self.get_response_format(self.model)
 
         system_message = self.get_system_message(
@@ -210,6 +214,7 @@ class SessionSummaryManager:
     ) -> Optional[SessionSummary]:
         """Creates a summary of the session"""
         log_debug("Creating session summary", center=True)
+        self.model = get_model(self.model)
         if self.model is None:
             return None
 
@@ -237,6 +242,7 @@ class SessionSummaryManager:
     ) -> Optional[SessionSummary]:
         """Creates a summary of the session"""
         log_debug("Creating session summary", center=True)
+        self.model = get_model(self.model)
         if self.model is None:
             return None
 

agno/session/workflow.py

@@ -129,12 +129,12 @@ class WorkflowSession:
         context_parts = ["<workflow_history_context>"]
 
         for i, (input_str, response_str) in enumerate(history_data, 1):
-            context_parts.append(f"[run-{i}]")
+            context_parts.append(f"[Workflow Run-{i}]")
 
             if input_str:
-                context_parts.append(f"input: {input_str}")
+                context_parts.append(f"User input: {input_str}")
             if response_str:
-                context_parts.append(f"response: {response_str}")
+                context_parts.append(f"Workflow output: {response_str}")
 
             context_parts.append("")  # Empty line between runs
 
@@ -154,6 +154,7 @@ class WorkflowSession:
                     runs_data.append(run.to_dict())
                 except Exception as e:
                     raise ValueError(f"Serialization failed: {str(e)}")
+
         return {
             "session_id": self.session_id,
             "user_id": self.user_id,

agno/team/team.py

@@ -45,6 +45,7 @@ from agno.models.base import Model
 from agno.models.message import Message, MessageReferences
 from agno.models.metrics import Metrics
 from agno.models.response import ModelResponse, ModelResponseEvent
+from agno.models.utils import get_model
 from agno.reasoning.step import NextAction, ReasoningStep, ReasoningSteps
 from agno.run.agent import RunEvent, RunOutput, RunOutputEvent
 from agno.run.base import RunStatus
@@ -437,7 +438,7 @@ class Team:
         self,
         members: List[Union[Agent, "Team"]],
         id: Optional[str] = None,
-        model: Optional[Model] = None,
+        model: Optional[Union[Model, str]] = None,
         name: Optional[str] = None,
         role: Optional[str] = None,
         respond_directly: bool = False,
@@ -498,9 +499,9 @@ class Team:
         post_hooks: Optional[Union[List[Callable[..., Any]], List[BaseGuardrail]]] = None,
         input_schema: Optional[Type[BaseModel]] = None,
         output_schema: Optional[Type[BaseModel]] = None,
-        parser_model: Optional[Model] = None,
+        parser_model: Optional[Union[Model, str]] = None,
         parser_model_prompt: Optional[str] = None,
-        output_model: Optional[Model] = None,
+        output_model: Optional[Union[Model, str]] = None,
         output_model_prompt: Optional[str] = None,
         use_json_mode: bool = False,
         parse_response: bool = True,
@@ -514,7 +515,7 @@ class Team:
         add_session_summary_to_context: Optional[bool] = None,
         metadata: Optional[Dict[str, Any]] = None,
         reasoning: bool = False,
-        reasoning_model: Optional[Model] = None,
+        reasoning_model: Optional[Union[Model, str]] = None,
         reasoning_agent: Optional[Agent] = None,
         reasoning_min_steps: int = 1,
         reasoning_max_steps: int = 10,
@@ -535,7 +536,7 @@ class Team:
     ):
         self.members = members
 
-        self.model = model
+        self.model = model  # type: ignore[assignment]
 
         self.name = name
         self.id = id
@@ -618,9 +619,9 @@ class Team:
 
         self.input_schema = input_schema
         self.output_schema = output_schema
-        self.parser_model = parser_model
+        self.parser_model = parser_model  # type: ignore[assignment]
         self.parser_model_prompt = parser_model_prompt
-        self.output_model = output_model
+        self.output_model = output_model  # type: ignore[assignment]
         self.output_model_prompt = output_model_prompt
         self.use_json_mode = use_json_mode
         self.parse_response = parse_response
@@ -637,7 +638,7 @@ class Team:
         self.metadata = metadata
 
         self.reasoning = reasoning
-        self.reasoning_model = reasoning_model
+        self.reasoning_model = reasoning_model  # type: ignore[assignment]
         self.reasoning_agent = reasoning_agent
         self.reasoning_min_steps = reasoning_min_steps
         self.reasoning_max_steps = reasoning_max_steps
@@ -694,6 +695,8 @@ class Team:
         # Lazy-initialized shared thread pool executor for background tasks (memory, cultural knowledge, etc.)
         self._background_executor: Optional[Any] = None
 
+        self._resolve_models()
+
     @property
     def background_executor(self) -> Any:
         """Lazy initialization of shared thread pool executor for background tasks.
@@ -902,6 +905,17 @@ class Team:
         """Return True if the db the team is equipped with is an Async implementation"""
         return self.db is not None and isinstance(self.db, AsyncBaseDb)
 
+    def _resolve_models(self) -> None:
+        """Resolve model strings to Model instances."""
+        if self.model is not None:
+            self.model = get_model(self.model)
+        if self.reasoning_model is not None:
+            self.reasoning_model = get_model(self.reasoning_model)
+        if self.parser_model is not None:
+            self.parser_model = get_model(self.parser_model)
+        if self.output_model is not None:
+            self.output_model = get_model(self.output_model)
+
     def initialize_team(self, debug_mode: Optional[bool] = None) -> None:
         # Make sure for the team, we are using the team logger
         use_team_logger()
@@ -1593,6 +1607,7 @@ class Team:
                     tools=_tools,
                     response_format=response_format,
                     stream_events=stream_events,
+                    session_state=session_state,
                 ):
                     raise_if_cancelled(run_response.run_id)  # type: ignore
                     yield event
@@ -1604,6 +1619,7 @@ class Team:
                     tools=_tools,
                     response_format=response_format,
                     stream_events=stream_events,
+                    session_state=session_state,
                 ):
                     raise_if_cancelled(run_response.run_id)  # type: ignore
                     from agno.run.team import IntermediateRunContentEvent, RunContentEvent
@@ -1932,6 +1948,7 @@ class Team:
             team_id=self.id,
             team_name=self.name,
             metadata=metadata,
+            session_state=session_state,
             input=run_input,
         )
 
@@ -2439,6 +2456,7 @@ class Team:
                     tools=_tools,
                     response_format=response_format,
                     stream_events=stream_events,
+                    session_state=session_state,
                 ):
                     raise_if_cancelled(run_response.run_id)  # type: ignore
                     yield event
@@ -2450,6 +2468,7 @@ class Team:
                     tools=_tools,
                     response_format=response_format,
                     stream_events=stream_events,
+                    session_state=session_state,
                 ):
                     raise_if_cancelled(run_response.run_id)  # type: ignore
                     from agno.run.team import IntermediateRunContentEvent, RunContentEvent
@@ -2771,6 +2790,7 @@ class Team:
             team_id=self.id,
             team_name=self.name,
             metadata=metadata,
+            session_state=session_state,
             input=run_input,
         )
 
@@ -2930,6 +2950,12 @@ class Team:
         if model_response.audio is not None:
             run_response.response_audio = model_response.audio
 
+        # Update session_state with changes from model response
+        if model_response.updated_session_state is not None and run_response.session_state is not None:
+            from agno.utils.merge_dict import merge_dictionaries
+
+            merge_dictionaries(run_response.session_state, model_response.updated_session_state)
+
         # Build a list of messages that should be added to the RunOutput
         messages_for_run_response = [m for m in run_messages.messages if m.add_to_agent_memory]
 
@@ -2954,6 +2980,7 @@ class Team:
         tools: Optional[List[Union[Function, dict]]] = None,
         response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
         stream_events: bool = False,
+        session_state: Optional[Dict[str, Any]] = None,
     ) -> Iterator[Union[TeamRunOutputEvent, RunOutputEvent]]:
         self.model = cast(Model, self.model)
 
@@ -2985,6 +3012,7 @@ class Team:
                 reasoning_state=reasoning_state,
                 stream_events=stream_events,
                 parse_structured_output=self.should_parse_structured_output,
+                session_state=session_state,
             )
 
         # 3. Update TeamRunOutput
@@ -3036,6 +3064,7 @@ class Team:
         tools: Optional[List[Union[Function, dict]]] = None,
         response_format: Optional[Union[Dict, Type[BaseModel]]] = None,
         stream_events: bool = False,
+        session_state: Optional[Dict[str, Any]] = None,
     ) -> AsyncIterator[Union[TeamRunOutputEvent, RunOutputEvent]]:
         self.model = cast(Model, self.model)
 
@@ -3068,6 +3097,7 @@ class Team:
                 reasoning_state=reasoning_state,
                 stream_events=stream_events,
                 parse_structured_output=self.should_parse_structured_output,
+                session_state=session_state,
             ):
                 yield event
 
@@ -3122,6 +3152,7 @@ class Team:
         reasoning_state: Optional[Dict[str, Any]] = None,
         stream_events: bool = False,
         parse_structured_output: bool = False,
+        session_state: Optional[Dict[str, Any]] = None,
     ) -> Iterator[Union[TeamRunOutputEvent, RunOutputEvent]]:
         if isinstance(model_response_event, tuple(get_args(RunOutputEvent))) or isinstance(
             model_response_event, tuple(get_args(TeamRunOutputEvent))
@@ -3299,10 +3330,15 @@ class Team:
 
             # If the model response is a tool_call_completed, update the existing tool call in the run_response
             elif model_response_event.event == ModelResponseEvent.tool_call_completed.value:
-                if model_response_event.updated_session_state is not None and session.session_data is not None:
-                    merge_dictionaries(
-                        session.session_data["session_state"], model_response_event.updated_session_state
-                    )
+                if model_response_event.updated_session_state is not None:
+                    # Update the session_state variable that TeamRunOutput references
+                    if session_state is not None:
+                        merge_dictionaries(session_state, model_response_event.updated_session_state)
+                    # Also update the DB session object
+                    if session.session_data is not None:
+                        merge_dictionaries(
+                            session.session_data["session_state"], model_response_event.updated_session_state
+                        )
 
                 if model_response_event.images is not None:
                     for image in model_response_event.images:
@@ -4498,7 +4534,7 @@ class Team:
                         store_events=self.store_events,
                     )
             else:
-                log_warning(
+                log_info(
                     f"Reasoning model: {reasoning_model.__class__.__name__} is not a native reasoning model, defaulting to manual Chain-of-Thought reasoning"
                 )
                 use_default_reasoning = True
@@ -4779,7 +4815,7 @@ class Team:
                         store_events=self.store_events,
                     )
             else:
-                log_warning(
+                log_info(
                     f"Reasoning model: {reasoning_model.__class__.__name__} is not a native reasoning model, defaulting to manual Chain-of-Thought reasoning"
                 )
                 use_default_reasoning = True