feat(callbacks): Callbacks can edit messages

docs/agents/callbacks.md

@@ -54,14 +54,42 @@ property so that callbacks can access information in a framework-agnostic way.
 
 You can see what attributes are available for LLM Calls and Tool Executions by examining the [`GenAI`][any_agent.tracing.attributes.GenAI] class.
 
+### Framework State
+
+In addition to the span attributes, callbacks can access and modify framework-specific objects through [`Context.framework_state`][any_agent.callbacks.context.Context.framework_state].
+
+This allows callbacks to directly manipulate the agent's execution, such as:
+
+- Modifying messages before they're sent to the LLM
+- Modifying the LLM's response after generation
+- Injecting prompts mid-execution
+- Changing user queries dynamically
+
+#### Helper Methods
+
+The `framework_state` provides helper methods to work with messages in a normalized format:
+
+**`get_messages()`**: Get messages as a list of dicts with `role` and `content` keys
+
+**`set_messages()`**: Set messages from a list of dicts with `role` and `content` keys
+
+These methods handle framework-specific message formats internally, providing a consistent API across frameworks.
+
+!!! note "Availability"
+
+    The `get_messages()` and `set_messages()` methods are **only available in `before_llm_call` callbacks**.
+
+    - In `before_llm_call`: You can read and modify the messages that will be sent to the LLM
+    - In other callbacks (`after_llm_call`, `before_tool_execution`, `after_tool_execution`, etc.): These methods will raise `NotImplementedError`
+
 ## Implementing Callbacks
 
 All callbacks must inherit from the base [`Callback`][any_agent.callbacks.base.Callback] class and can choose to implement any subset of the available callback methods. These methods include:
 
 | Callback Method | Description |
 |:----------------:|:------------:|
 | before_agent_invocation | Should be used to check the Context before the agent is invoked. |
-| befor_llm_call | Should be used before the chat history hits the LLM. |
+| before_llm_call | Should be used before the chat history hits the LLM. |
 | after_llm_call | Should be used once LLM output is generated, before it appends to the chat history. |
 | before_tool_execution | Should be used to check the Context before tool execution. |
 | after_tool_execution | Should be used once tool has been executed, before the output is appended to chat history. |
@@ -136,7 +164,7 @@ Callbacks are provided to the agent using the [`AgentConfig.callbacks`][any_agen
     agent = AnyAgent.create(
         "tinyagent",
         AgentConfig(
-            model_id="gpt-4.1-nano",
+            model_id="openai:gpt-4.1-nano",
             instructions="Use the tools to find an answer",
             tools=[search_web, visit_webpage],
             callbacks=[
@@ -157,7 +185,7 @@ Callbacks are provided to the agent using the [`AgentConfig.callbacks`][any_agen
     agent = AnyAgent.create(
         "tinyagent",
         AgentConfig(
-            model_id="gpt-4.1-nano",
+            model_id="openai:gpt-4.1-nano",
             instructions="Use the tools to find an answer",
             tools=[search_web, visit_webpage],
             callbacks=[
@@ -272,3 +300,50 @@ class LimitToolExecutions(Callback):
 
         return context
 ```
+
+## Example: Modifying prompts dynamically
+
+You can use callbacks to modify the prompt being sent to the LLM. This is useful for injecting instructions or reminders mid-execution:
+
+```python
+from any_agent.callbacks.base import Callback
+from any_agent.callbacks.context import Context
+
+class InjectReminderCallback(Callback):
+    def __init__(self, reminder: str, every_n_calls: int = 5):
+        self.reminder = reminder
+        self.every_n_calls = every_n_calls
+        self.call_count = 0
+
+    def before_llm_call(self, context: Context, *args, **kwargs) -> Context:
+        self.call_count += 1
+
+        if self.call_count % self.every_n_calls == 0:
+            try:
+                messages = context.framework_state.get_messages()
+                if messages:
+                    messages[-1]["content"] += f"\n\n{self.reminder}"
+                    context.framework_state.set_messages(messages)
+            except NotImplementedError:
+                pass
+
+        return context
+```
+
+Example usage:
+
+```python
+from any_agent import AgentConfig, AnyAgent
+
+callback = InjectReminderCallback(
+    reminder="Remember to use the Todo tool to track your tasks!",
+    every_n_calls=5
+)
+
+config = AgentConfig(
+    model_id="openai:gpt-4o-mini",
+    instructions="You are a helpful assistant.",
+    callbacks=[callback],
+)
+# ... Continue to create and run agent
+```

docs/api/callbacks.md

@@ -4,6 +4,8 @@
 
 ::: any_agent.callbacks.context.Context
 
+::: any_agent.callbacks.context.FrameworkState
+
 ::: any_agent.callbacks.span_print.ConsolePrintSpan
 
 ::: any_agent.callbacks.get_default_callbacks

src/any_agent/callbacks/__init__.py

@@ -1,8 +1,8 @@
 from .base import Callback
-from .context import Context
+from .context import Context, FrameworkState
 from .span_print import ConsolePrintSpan
 
-__all__ = ["Callback", "ConsolePrintSpan", "Context"]
+__all__ = ["Callback", "ConsolePrintSpan", "Context", "FrameworkState"]
 
 
 def get_default_callbacks() -> list[Callback]:

src/any_agent/callbacks/context.py

@@ -1,14 +1,87 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from opentelemetry.trace import Span, Tracer
 
     from any_agent.tracing.agent_trace import AgentTrace
 
 
+@dataclass
+class FrameworkState:
+    """Framework-specific state that can be accessed and modified by callbacks.
+
+    This object provides a consistent interface for accessing framework state across
+    different agent frameworks, while the actual content is framework-specific.
+    """
+
+    messages: list[dict[str, Any]] = field(default_factory=list)
+    """Internal storage for messages. Use get_messages() and set_messages() instead."""
+
+    _message_getter: Callable[[], list[dict[str, Any]]] | None = field(
+        default=None, repr=False
+    )
+    """Framework-specific message getter function."""
+
+    _message_setter: Callable[[list[dict[str, Any]]], None] | None = field(
+        default=None, repr=False
+    )
+    """Framework-specific message setter function."""
+
+    def get_messages(self) -> list[dict[str, Any]]:
+        """Get messages in a normalized dict format.
+
+        Returns a list of message dicts with 'role' and 'content' keys.
+        Works consistently across all frameworks.
+
+        Returns:
+            List of message dicts with 'role' and 'content' keys.
+
+        Raises:
+            NotImplementedError: If the framework doesn't support message access yet.
+
+        Example:
+            ```python
+            messages = context.framework_state.get_messages()
+            # [{"role": "user", "content": "Hello"}]
+            ```
+
+        """
+        if self._message_getter is None:
+            msg = "get_messages() is not implemented for this framework yet"
+            raise NotImplementedError(msg)
+        return self._message_getter()
+
+    def set_messages(self, messages: list[dict[str, Any]]) -> None:
+        """Set messages from a normalized dict format.
+
+        Accepts a list of message dicts with 'role' and 'content' keys and
+        converts them to the framework-specific format.
+
+        Args:
+            messages: List of message dicts with 'role' and 'content' keys.
+
+        Raises:
+            NotImplementedError: If the framework doesn't support message modification yet.
+
+        Example:
+            ```python
+            messages = context.framework_state.get_messages()
+            messages[-1]["content"] = "Say hello"
+            context.framework_state.set_messages(messages)
+            ```
+
+        """
+        if self._message_setter is None:
+            msg = "set_messages() is not implemented for this framework yet"
+            raise NotImplementedError(msg)
+        self._message_setter(messages)
+
+
 @dataclass
 class Context:
     """Object that will be shared across callbacks.
@@ -31,3 +104,20 @@ class Context:
 
     shared: dict[str, Any]
     """Can be used to store arbitrary information for sharing across callbacks."""
+
+    framework_state: FrameworkState
+    """Framework-specific state that can be accessed and modified by callbacks.
+
+    Provides consistent access to framework state across different agent frameworks.
+    See [`FrameworkState`][any_agent.callbacks.context.FrameworkState] for available attributes.
+
+    Example:
+        ```python
+        class ModifyPromptCallback(Callback):
+            def before_llm_call(self, context: Context, *args, **kwargs) -> Context:
+                # Modify the last message content
+                if context.framework_state.messages:
+                    context.framework_state.messages[-1]["content"] = "Say hello"
+                return context
+        ```
+    """

src/any_agent/callbacks/span_generation/agno.py

@@ -23,26 +23,62 @@ def before_llm_call(self, context: Context, *args, **kwargs):
 
     def after_llm_call(self, context: Context, *args, **kwargs) -> Context:
         output: str | list[dict[str, Any]] = ""
+        input_tokens: int = 0
+        output_tokens: int = 0
+
         if assistant_message := kwargs.get("assistant_message"):
-            if content := getattr(assistant_message, "content", None):
-                output = str(content)
-            if tool_calls := getattr(assistant_message, "tool_calls", None):
-                output = [
-                    {
-                        "tool.name": tool.get("function", {}).get("name", "No name"),
-                        "tool.args": tool.get("function", {}).get(
-                            "arguments", "No args"
-                        ),
-                    }
-                    for tool in tool_calls
-                ]
+            if hasattr(assistant_message, "choices"):
+                choices = getattr(assistant_message, "choices", [])
+                if choices and len(choices) > 0:
+                    choice = choices[0]
+                    message = getattr(choice, "message", None)
+                    if message:
+                        if content := getattr(message, "content", None):
+                            output = str(content)
+                        if tool_calls := getattr(message, "tool_calls", None):
+                            output = [
+                                {
+                                    "tool.name": getattr(
+                                        getattr(tool, "function", None),
+                                        "name",
+                                        "No name",
+                                    ),
+                                    "tool.args": getattr(
+                                        getattr(tool, "function", None),
+                                        "arguments",
+                                        "No args",
+                                    ),
+                                }
+                                for tool in tool_calls
+                            ]
+
+                if usage := getattr(assistant_message, "usage", None):
+                    input_tokens = getattr(usage, "input_tokens", 0) or getattr(
+                        usage, "prompt_tokens", 0
+                    )
+                    output_tokens = getattr(usage, "output_tokens", 0) or getattr(
+                        usage, "completion_tokens", 0
+                    )
+            else:
+                if content := getattr(assistant_message, "content", None):
+                    output = str(content)
+                if tool_calls := getattr(assistant_message, "tool_calls", None):
+                    output = [
+                        {
+                            "tool.name": tool.get("function", {}).get(
+                                "name", "No name"
+                            ),
+                            "tool.args": tool.get("function", {}).get(
+                                "arguments", "No args"
+                            ),
+                        }
+                        for tool in tool_calls
+                    ]
 
-            metrics: MessageMetrics | None
-            input_tokens: int = 0
-            output_tokens: int = 0
-            if metrics := getattr(assistant_message, "metrics", None):
-                input_tokens = metrics.input_tokens
-                output_tokens = metrics.output_tokens
+                metrics: MessageMetrics | None
+                if metrics := getattr(assistant_message, "metrics", None):
+                    input_tokens = metrics.input_tokens
+                    output_tokens = metrics.output_tokens
 
             context = self._set_llm_output(context, output, input_tokens, output_tokens)