agno 2.2.13__py3-none-any.whl → 2.4.3__py3-none-any.whl

agno/tools/file_generation.py

@@ -108,14 +108,16 @@ class FileGenerationTools(Toolkit):
             # Save file to disk (if output_directory is set)
             file_path = self._save_file_to_disk(json_content, filename)
 
+            content_bytes = json_content.encode("utf-8")
+
             # Create FileArtifact
             file_artifact = File(
                 id=str(uuid4()),
-                content=json_content,
+                content=content_bytes,
                 mime_type="application/json",
                 file_type="json",
                 filename=filename,
-                size=len(json_content.encode("utf-8")),
+                size=len(content_bytes),
                 filepath=file_path if file_path else None,
             )
 
@@ -195,14 +197,16 @@ class FileGenerationTools(Toolkit):
             # Save file to disk (if output_directory is set)
             file_path = self._save_file_to_disk(csv_content, filename)
 
+            content_bytes = csv_content.encode("utf-8")
+
             # Create FileArtifact
             file_artifact = File(
                 id=str(uuid4()),
-                content=csv_content,
+                content=content_bytes,
                 mime_type="text/csv",
                 file_type="csv",
                 filename=filename,
-                size=len(csv_content.encode("utf-8")),
+                size=len(content_bytes),
                 filepath=file_path if file_path else None,
             )
 
@@ -325,14 +329,16 @@ class FileGenerationTools(Toolkit):
             # Save file to disk (if output_directory is set)
             file_path = self._save_file_to_disk(content, filename)
 
+            content_bytes = content.encode("utf-8")
+
             # Create FileArtifact
             file_artifact = File(
                 id=str(uuid4()),
-                content=content,
+                content=content_bytes,
                 mime_type="text/plain",
                 file_type="txt",
                 filename=filename,
-                size=len(content.encode("utf-8")),
+                size=len(content_bytes),
                 filepath=file_path if file_path else None,
             )
 

agno/tools/firecrawl.py

@@ -101,8 +101,10 @@ class FirecrawlTools(Toolkit):
             The results of the crawling.
         """
         params: Dict[str, Any] = {}
-        if self.limit or limit:
-            params["limit"] = self.limit or limit
+        if self.limit is not None:
+            params["limit"] = self.limit
+        elif limit is not None:
+            params["limit"] = limit
         if self.formats:
             params["scrape_options"] = ScrapeOptions(formats=self.formats)  # type: ignore
 
@@ -129,15 +131,21 @@ class FirecrawlTools(Toolkit):
             limit (int): The maximum number of results to return.
         """
         params: Dict[str, Any] = {}
-        if self.limit or limit:
-            params["limit"] = self.limit or limit
+        if self.limit is not None:
+            params["limit"] = self.limit
+        elif limit is not None:
+            params["limit"] = limit
         if self.formats:
             params["scrape_options"] = ScrapeOptions(formats=self.formats)  # type: ignore
         if self.search_params:
             params.update(self.search_params)
 
         search_result = self.app.search(query, **params)
-        if search_result.success:
-            return json.dumps(search_result.data, cls=CustomJSONEncoder)
+
+        if hasattr(search_result, "success"):
+            if search_result.success:
+                return json.dumps(search_result.data, cls=CustomJSONEncoder)
+            else:
+                return f"Error searching with the Firecrawl tool: {search_result.error}"
         else:
-            return "Error searching with the Firecrawl tool: " + search_result.error
+            return json.dumps(search_result.model_dump(), cls=CustomJSONEncoder)

agno/tools/function.py

@@ -54,9 +54,17 @@ class UserInputField:
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "UserInputField":
+        type_mapping = {"str": str, "int": int, "float": float, "bool": bool, "list": list, "dict": dict}
+        field_type_raw = data["field_type"]
+        if isinstance(field_type_raw, str):
+            field_type = type_mapping.get(field_type_raw, str)
+        elif isinstance(field_type_raw, type):
+            field_type = field_type_raw
+        else:
+            field_type = str
         return cls(
             name=data["name"],
-            field_type=eval(data["field_type"]),  # Convert string type name to actual type
+            field_type=field_type,
             description=data["description"],
             value=data["value"],
         )
@@ -125,10 +133,6 @@ class Function(BaseModel):
     _team: Optional[Any] = None
     # The run context that the function is associated with
     _run_context: Optional[RunContext] = None
-    # The session state that the function is associated with
-    _session_state: Optional[Dict[str, Any]] = None
-    # The dependencies that the function is associated with
-    _dependencies: Optional[Dict[str, Any]] = None
 
     # Media context that the function is associated with
     _images: Optional[Sequence[Image]] = None
@@ -142,6 +146,19 @@ class Function(BaseModel):
             include={"name", "description", "parameters", "strict", "requires_confirmation", "external_execution"},
         )
 
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Function":
+        """Reconstruct a Function from a dictionary."""
+
+        return cls(
+            name=data.get("name"),
+            description=data.get("description"),
+            parameters=data.get("parameters"),
+            strict=data.get("strict"),
+            requires_confirmation=data.get("requires_confirmation", False),
+            external_execution=data.get("external_execution", False),
+        )
+
     def model_copy(self, *, deep: bool = False) -> "Function":
         """
         Override model_copy to handle callable fields that can't be deep copied (pickled).
@@ -201,10 +218,6 @@ class Function(BaseModel):
                 del type_hints["team"]
             if "run_context" in sig.parameters and "run_context" in type_hints:
                 del type_hints["run_context"]
-            if "session_state" in sig.parameters and "session_state" in type_hints:
-                del type_hints["session_state"]
-            if "dependencies" in sig.parameters and "dependencies" in type_hints:
-                del type_hints["dependencies"]
 
             # Remove media parameters from type hints as they are injected automatically
             if "images" in sig.parameters and "images" in type_hints:
@@ -227,8 +240,6 @@ class Function(BaseModel):
                     "agent",
                     "team",
                     "run_context",
-                    "session_state",
-                    "dependencies",
                     "self",
                     "images",
                     "videos",
@@ -268,8 +279,6 @@ class Function(BaseModel):
                         "agent",
                         "team",
                         "run_context",
-                        "session_state",
-                        "dependencies",
                         "self",
                         "images",
                         "videos",
@@ -288,8 +297,6 @@ class Function(BaseModel):
                         "agent",
                         "team",
                         "run_context",
-                        "session_state",
-                        "dependencies",
                         "self",
                         "images",
                         "videos",
@@ -346,10 +353,6 @@ class Function(BaseModel):
                 del type_hints["team"]
             if "run_context" in sig.parameters and "run_context" in type_hints:
                 del type_hints["run_context"]
-            if "session_state" in sig.parameters and "session_state" in type_hints:
-                del type_hints["session_state"]
-            if "dependencies" in sig.parameters and "dependencies" in type_hints:
-                del type_hints["dependencies"]
             if "images" in sig.parameters and "images" in type_hints:
                 del type_hints["images"]
             if "videos" in sig.parameters and "videos" in type_hints:
@@ -366,8 +369,6 @@ class Function(BaseModel):
                 "agent",
                 "team",
                 "run_context",
-                "session_state",
-                "dependencies",
                 "self",
                 "images",
                 "videos",
@@ -480,10 +481,15 @@ class Function(BaseModel):
         # Don't wrap callables that are already wrapped with validate_call
         elif getattr(func, "_wrapped_for_validation", False):
             return func
-        # Don't wrap functions with session_state parameter
-        # session_state needs to be passed by reference, not copied by pydantic's validation
-        elif "session_state" in signature(func).parameters:
+
+        # Don't wrap functions with framework-injected parameters
+        # These parameters (agent, team) are
+        # injected by the framework at runtime and shouldn't be validated by Pydantic
+        sig = signature(func)
+        framework_params = {"agent", "team"}
+        if framework_params & set(sig.parameters.keys()):
             return func
+
         # Wrap the callable with validate_call
         else:
             wrapped = validate_call(func, config=dict(arbitrary_types_allowed=True))  # type: ignore
@@ -538,8 +544,6 @@ class Function(BaseModel):
                 "agent",
                 "team",
                 "run_context",
-                "session_state",
-                "dependencies",
                 "images",
                 "videos",
                 "audios",
@@ -561,10 +565,6 @@ class Function(BaseModel):
             del copy_entrypoint_args["team"]
         if "run_context" in copy_entrypoint_args:
             del copy_entrypoint_args["run_context"]
-        if "session_state" in copy_entrypoint_args:
-            del copy_entrypoint_args["session_state"]
-        if "dependencies" in copy_entrypoint_args:
-            del copy_entrypoint_args["dependencies"]
         if "images" in copy_entrypoint_args:
             del copy_entrypoint_args["images"]
         if "videos" in copy_entrypoint_args:
@@ -691,21 +691,15 @@ class FunctionCall(BaseModel):
                 from inspect import signature
 
                 pre_hook_args = {}
-                # Check if the pre-hook has and agent argument
+                # Check if the pre-hook has an agent argument
                 if "agent" in signature(self.function.pre_hook).parameters:
                     pre_hook_args["agent"] = self.function._agent
-                # Check if the pre-hook has an team argument
+                # Check if the pre-hook has a team argument
                 if "team" in signature(self.function.pre_hook).parameters:
                     pre_hook_args["team"] = self.function._team
-                # Check if the pre-hook has an session_state argument
+                # Check if the pre-hook has a run_context argument
                 if "run_context" in signature(self.function.pre_hook).parameters:
                     pre_hook_args["run_context"] = self.function._run_context
-                # Check if the pre-hook has an session_state argument
-                if "session_state" in signature(self.function.pre_hook).parameters:
-                    pre_hook_args["session_state"] = self.function._session_state
-                # Check if the pre-hook has an dependencies argument
-                if "dependencies" in signature(self.function.pre_hook).parameters:
-                    pre_hook_args["dependencies"] = self.function._dependencies
                 # Check if the pre-hook has an fc argument
                 if "fc" in signature(self.function.pre_hook).parameters:
                     pre_hook_args["fc"] = self
@@ -725,21 +719,15 @@ class FunctionCall(BaseModel):
                 from inspect import signature
 
                 post_hook_args = {}
-                # Check if the post-hook has and agent argument
+                # Check if the post-hook has an agent argument
                 if "agent" in signature(self.function.post_hook).parameters:
                     post_hook_args["agent"] = self.function._agent
-                # Check if the post-hook has an team argument
+                # Check if the post-hook has a team argument
                 if "team" in signature(self.function.post_hook).parameters:
                     post_hook_args["team"] = self.function._team
-                # Check if the post-hook has an session_state argument
+                # Check if the post-hook has a run_context argument
                 if "run_context" in signature(self.function.post_hook).parameters:
                     post_hook_args["run_context"] = self.function._run_context
-                # Check if the post-hook has an session_state argument
-                if "session_state" in signature(self.function.post_hook).parameters:
-                    post_hook_args["session_state"] = self.function._session_state
-                # Check if the post-hook has an dependencies argument
-                if "dependencies" in signature(self.function.post_hook).parameters:
-                    post_hook_args["dependencies"] = self.function._dependencies
                 # Check if the post-hook has an fc argument
                 if "fc" in signature(self.function.post_hook).parameters:
                     post_hook_args["fc"] = self
@@ -760,18 +748,12 @@ class FunctionCall(BaseModel):
         # Check if the entrypoint has an agent argument
         if "agent" in signature(self.function.entrypoint).parameters:  # type: ignore
             entrypoint_args["agent"] = self.function._agent
-        # Check if the entrypoint has an team argument
+        # Check if the entrypoint has a team argument
         if "team" in signature(self.function.entrypoint).parameters:  # type: ignore
             entrypoint_args["team"] = self.function._team
-        # Check if the entrypoint has an run_context argument
+        # Check if the entrypoint has a run_context argument
         if "run_context" in signature(self.function.entrypoint).parameters:  # type: ignore
             entrypoint_args["run_context"] = self.function._run_context
-        # Check if the entrypoint has an session_state argument
-        if "session_state" in signature(self.function.entrypoint).parameters:  # type: ignore
-            entrypoint_args["session_state"] = self.function._session_state
-        # Check if the entrypoint has an dependencies argument
-        if "dependencies" in signature(self.function.entrypoint).parameters:  # type: ignore
-            entrypoint_args["dependencies"] = self.function._dependencies
         # Check if the entrypoint has an fc argument
         if "fc" in signature(self.function.entrypoint).parameters:  # type: ignore
             entrypoint_args["fc"] = self
@@ -795,18 +777,12 @@ class FunctionCall(BaseModel):
         # Check if the hook has an agent argument
         if "agent" in signature(hook).parameters:
             hook_args["agent"] = self.function._agent
-        # Check if the hook has an team argument
+        # Check if the hook has a team argument
         if "team" in signature(hook).parameters:
             hook_args["team"] = self.function._team
-        # Check if the hook has an run_context argument
+        # Check if the hook has a run_context argument
         if "run_context" in signature(hook).parameters:
             hook_args["run_context"] = self.function._run_context
-        # Check if the hook has an session_state argument
-        if "session_state" in signature(hook).parameters:
-            hook_args["session_state"] = self.function._session_state
-        # Check if the hook has an dependencies argument
-        if "dependencies" in signature(hook).parameters:
-            hook_args["dependencies"] = self.function._dependencies
         if "name" in signature(hook).parameters:
             hook_args["name"] = name
         if "function_name" in signature(hook).parameters:
@@ -897,7 +873,7 @@ class FunctionCall(BaseModel):
                 return FunctionExecutionResult(status="success", result=cached_result)
 
         # Execute function
-        execution_result = None
+        execution_result: FunctionExecutionResult
         exception_to_raise = None
 
         try:
@@ -908,21 +884,16 @@ class FunctionCall(BaseModel):
             else:
                 result = self.function.entrypoint(**entrypoint_args, **self.arguments)  # type: ignore
 
-            updated_session_state = None
-            if entrypoint_args.get("run_context") is not None:
-                run_context = entrypoint_args.get("run_context")
-                updated_session_state = (
-                    run_context.session_state
-                    if run_context is not None and run_context.session_state is not None
-                    else None
-                )
-            else:
-                if self.function._session_state is not None:
-                    updated_session_state = self.function._session_state
-
             # Handle generator case
             if isgenerator(result):
                 self.result = result  # Store generator directly, can't cache
+                # For generators, don't capture updated_session_state yet -
+                # session_state is passed by reference, so mutations made during
+                # generator iteration are already reflected in the original dict.
+                # Returning None prevents stale state from being merged later.
+                execution_result = FunctionExecutionResult(
+                    status="success", result=self.result, updated_session_state=None
+                )
             else:
                 self.result = result
                 # Only cache non-generator results
@@ -931,14 +902,24 @@ class FunctionCall(BaseModel):
                     cache_file = self.function._get_cache_file_path(cache_key)
                     self.function._save_to_cache(cache_file, self.result)
 
-            execution_result = FunctionExecutionResult(
-                status="success", result=self.result, updated_session_state=updated_session_state
-            )
+                updated_session_state = None
+                if entrypoint_args.get("run_context") is not None:
+                    run_context = entrypoint_args.get("run_context")
+                    updated_session_state = (
+                        run_context.session_state
+                        if run_context is not None and run_context.session_state is not None
+                        else None
+                    )
+
+                execution_result = FunctionExecutionResult(
+                    status="success", result=self.result, updated_session_state=updated_session_state
+                )
 
         except AgentRunException as e:
             log_debug(f"{e.__class__.__name__}: {e}")
             self.error = str(e)
             exception_to_raise = e
+            execution_result = FunctionExecutionResult(status="failure", error=str(e))
         except Exception as e:
             log_warning(f"Could not run function {self.get_call_str()}")
             log_exception(e)
@@ -948,10 +929,10 @@ class FunctionCall(BaseModel):
         finally:
             self._handle_post_hook()
 
-            if exception_to_raise is not None:
-                raise exception_to_raise
+        if exception_to_raise is not None:
+            raise exception_to_raise
 
-            return execution_result  # type: ignore[return-value]
+        return execution_result
 
     async def _handle_pre_hook_async(self):
         """Handles the async pre-hook for the function call."""
@@ -963,18 +944,12 @@ class FunctionCall(BaseModel):
                 # Check if the pre-hook has an agent argument
                 if "agent" in signature(self.function.pre_hook).parameters:
                     pre_hook_args["agent"] = self.function._agent
-                # Check if the pre-hook has an team argument
+                # Check if the pre-hook has a team argument
                 if "team" in signature(self.function.pre_hook).parameters:
                     pre_hook_args["team"] = self.function._team
-                # Check if the pre-hook has an run_context argument
+                # Check if the pre-hook has a run_context argument
                 if "run_context" in signature(self.function.pre_hook).parameters:
                     pre_hook_args["run_context"] = self.function._run_context
-                # Check if the pre-hook has an session_state argument
-                if "session_state" in signature(self.function.pre_hook).parameters:
-                    pre_hook_args["session_state"] = self.function._session_state
-                # Check if the pre-hook has an dependencies argument
-                if "dependencies" in signature(self.function.pre_hook).parameters:
-                    pre_hook_args["dependencies"] = self.function._dependencies
                 # Check if the pre-hook has an fc argument
                 if "fc" in signature(self.function.pre_hook).parameters:
                     pre_hook_args["fc"] = self
@@ -998,19 +973,12 @@ class FunctionCall(BaseModel):
                 # Check if the post-hook has an agent argument
                 if "agent" in signature(self.function.post_hook).parameters:
                     post_hook_args["agent"] = self.function._agent
-                # Check if the post-hook has an team argument
+                # Check if the post-hook has a team argument
                 if "team" in signature(self.function.post_hook).parameters:
                     post_hook_args["team"] = self.function._team
-                # Check if the post-hook has an run_context argument
+                # Check if the post-hook has a run_context argument
                 if "run_context" in signature(self.function.post_hook).parameters:
                     post_hook_args["run_context"] = self.function._run_context
-                # Check if the post-hook has an session_state argument
-                if "session_state" in signature(self.function.post_hook).parameters:
-                    post_hook_args["session_state"] = self.function._session_state
-                # Check if the post-hook has an dependencies argument
-                if "dependencies" in signature(self.function.post_hook).parameters:
-                    post_hook_args["dependencies"] = self.function._dependencies
-
                 # Check if the post-hook has an fc argument
                 if "fc" in signature(self.function.post_hook).parameters:
                     post_hook_args["fc"] = self
@@ -1117,7 +1085,7 @@ class FunctionCall(BaseModel):
                 return FunctionExecutionResult(status="success", result=cached_result)
 
         # Execute function
-        execution_result = None
+        execution_result: FunctionExecutionResult
         exception_to_raise = None
 
         try:
@@ -1131,10 +1099,15 @@ class FunctionCall(BaseModel):
                 else:
                     result = self.function.entrypoint(**entrypoint_args, **self.arguments)
 
+                # Handle both sync and async entrypoints
                 if isasyncgenfunction(self.function.entrypoint):
                     self.result = result  # Store async generator directly
+                elif iscoroutinefunction(self.function.entrypoint):
+                    self.result = await result  # Await coroutine result
+                elif isgeneratorfunction(self.function.entrypoint):
+                    self.result = result  # Store sync generator directly
                 else:
-                    self.result = await result
+                    self.result = result  # Sync function, result is already computed
 
             # Only cache if not a generator
             if self.function.cache_results and not (isgenerator(self.result) or isasyncgen(self.result)):
@@ -1142,14 +1115,21 @@ class FunctionCall(BaseModel):
                 cache_file = self.function._get_cache_file_path(cache_key)
                 self.function._save_to_cache(cache_file, self.result)
 
-            updated_session_state = None
-            if entrypoint_args.get("run_context") is not None:
-                run_context = entrypoint_args.get("run_context")
-                updated_session_state = (
-                    run_context.session_state
-                    if run_context is not None and run_context.session_state is not None
-                    else None
-                )
+            # For generators, don't capture updated_session_state -
+            # session_state is passed by reference, so mutations made during
+            # generator iteration are already reflected in the original dict.
+            # Returning None prevents stale state from being merged later.
+            if isgenerator(self.result) or isasyncgen(self.result):
+                updated_session_state = None
+            else:
+                updated_session_state = None
+                if entrypoint_args.get("run_context") is not None:
+                    run_context = entrypoint_args.get("run_context")
+                    updated_session_state = (
+                        run_context.session_state
+                        if run_context is not None and run_context.session_state is not None
+                        else None
+                    )
 
             execution_result = FunctionExecutionResult(
                 status="success", result=self.result, updated_session_state=updated_session_state
@@ -1159,6 +1139,7 @@ class FunctionCall(BaseModel):
             log_debug(f"{e.__class__.__name__}: {e}")
             self.error = str(e)
             exception_to_raise = e
+            execution_result = FunctionExecutionResult(status="failure", error=str(e))
         except Exception as e:
             log_warning(f"Could not run function {self.get_call_str()}")
             log_exception(e)
@@ -1171,10 +1152,10 @@ class FunctionCall(BaseModel):
             else:
                 self._handle_post_hook()
 
-            if exception_to_raise is not None:
-                raise exception_to_raise
+        if exception_to_raise is not None:
+            raise exception_to_raise
 
-            return execution_result  # type: ignore[return-value]
+        return execution_result
 
 
 class ToolResult(BaseModel):

agno/tools/google_bigquery.py

@@ -11,6 +11,15 @@ except ImportError:
     raise ImportError("`bigquery` not installed. Please install using `pip install google-cloud-bigquery`")
 
 
+def _clean_sql(sql: str) -> str:
+    """Clean SQL query by normalizing whitespace while preserving token boundaries.
+
+    Replaces newlines with spaces (not empty strings) to prevent line comments
+    from swallowing subsequent SQL statements.
+    """
+    return sql.replace("\\n", " ").replace("\n", " ")
+
+
 class GoogleBigQueryTools(Toolkit):
     def __init__(
         self,
@@ -106,12 +115,12 @@ class GoogleBigQueryTools(Toolkit):
         """
         try:
             log_debug(f"Running Google SQL |\n{sql}")
-            cleaned_query = sql.replace("\\n", " ").replace("\n", "").replace("\\", "")
+            cleaned_query = _clean_sql(sql)
             job_config = bigquery.QueryJobConfig(default_dataset=f"{self.project}.{self.dataset}")
             query_job = self.client.query(cleaned_query, job_config)
             results = query_job.result()
             results_str = str([dict(row) for row in results])
-            return results_str.replace("\\", "").replace("\n", "")
+            return results_str.replace("\n", " ")
         except Exception as e:
             logger.error(f"Error while executing SQL: {e}")
             return ""

agno/tools/google_drive.py

@@ -69,6 +69,7 @@ from pathlib import Path
 from typing import Any, List, Optional, Union
 
 from agno.tools import Toolkit
+from agno.utils.log import log_error
 
 try:
     from google.auth.transport.requests import Request
@@ -202,7 +203,7 @@ class GoogleDriveTools(Toolkit):
             items = results.get("files", [])
             return items
         except Exception as error:
-            print(f"Could not list files: {error}")
+            log_error(f"Could not list files: {error}")
             return []
 
     @authenticate
@@ -238,7 +239,7 @@ class GoogleDriveTools(Toolkit):
             )
             return uploaded_file
         except Exception as error:
-            print(f"Could not upload file '{file_path}': {error}")
+            log_error(f"Could not upload file '{file_path}': {error}")
             return None
 
     @authenticate
@@ -266,5 +267,5 @@ class GoogleDriveTools(Toolkit):
                     print(f"Download progress: {int(status.progress() * 100)}%.")
             return dest_path
         except Exception as error:
-            print(f"Could not download file '{file_id}': {error}")
+            log_error(f"Could not download file '{file_id}': {error}")
             return None

agno/tools/knowledge.py

@@ -1,9 +1,10 @@
 import json
 from textwrap import dedent
-from typing import Any, Dict, List, Optional
+from typing import Any, List, Optional
 
 from agno.knowledge.document import Document
 from agno.knowledge.knowledge import Knowledge
+from agno.run import RunContext
 from agno.tools import Toolkit
 from agno.utils.log import log_debug, log_error
 
@@ -55,7 +56,7 @@ class KnowledgeTools(Toolkit):
             **kwargs,
         )
 
-    def think(self, session_state: Dict[str, Any], thought: str) -> str:
+    def think(self, run_context: RunContext, thought: str) -> str:
         """Use this tool as a scratchpad to reason about the question, refine your approach, brainstorm search terms, or revise your plan.
 
         Call `Think` whenever you need to figure out what to do next, analyze the user's question, or plan your approach.
@@ -71,8 +72,10 @@ class KnowledgeTools(Toolkit):
             log_debug(f"Thought: {thought}")
 
             # Add the thought to the Agent state
+            session_state = run_context.session_state
             if session_state is None:
                 session_state = {}
+                run_context.session_state = session_state
             if "thoughts" not in session_state:
                 session_state["thoughts"] = []
             session_state["thoughts"].append(thought)
@@ -89,7 +92,7 @@ class KnowledgeTools(Toolkit):
             log_error(f"Error recording thought: {e}")
             return f"Error recording thought: {e}"
 
-    def search_knowledge(self, session_state: Dict[str, Any], query: str) -> str:
+    def search_knowledge(self, run_context: RunContext, query: str) -> str:
         """Use this tool to search the knowledge base for relevant information.
         After thinking through the question, use this tool as many times as needed to search for relevant information.
 
@@ -111,7 +114,7 @@ class KnowledgeTools(Toolkit):
             log_error(f"Error searching knowledge base: {e}")
             return f"Error searching knowledge base: {e}"
 
-    def analyze(self, session_state: Dict[str, Any], analysis: str) -> str:
+    def analyze(self, run_context: RunContext, analysis: str) -> str:
         """Use this tool to evaluate whether the returned documents are correct and sufficient.
         If not, go back to "Think" or "Search" with refined queries.
 
@@ -125,8 +128,10 @@ class KnowledgeTools(Toolkit):
             log_debug(f"Analysis: {analysis}")
 
             # Add the thought to the Agent state
+            session_state = run_context.session_state
             if session_state is None:
                 session_state = {}
+                run_context.session_state = session_state
             if "analysis" not in session_state:
                 session_state["analysis"] = []
             session_state["analysis"].append(analysis)