splitting ctx.linearize() into two more fitting functions:

1) `ctx.render_for_generation()` used in generate calls
2) `ctx.full_event_log()` used for views on the whole context

also: fixed and inherent bug in `m.chat` which wrongly added the chat turn to the contexts (as single events, not as turn).

Author: HendrikStrobelt

docs/examples/agents/react.py

@@ -103,7 +103,7 @@ def react(
     react_toolbox: ReactToolbox,
 ):
     assert m.ctx.is_chat_context, "ReACT requires a chat context."
-    test_ctx_lin = m.ctx.linearize()
+    test_ctx_lin = m.ctx.render_for_generation()
     assert test_ctx_lin is not None and len(test_ctx_lin) == 0, (
         "ReACT expects a fresh context."
     )

docs/examples/agents/react_instruct.py

@@ -101,7 +101,7 @@ def react(
     react_toolbox: ReactToolbox,
 ):
     assert m.ctx.is_chat_context, "ReACT requires a chat context."
-    test_ctx_lin = m.ctx.linearize()
+    test_ctx_lin = m.ctx.render_for_generation()
     assert test_ctx_lin is not None and len(test_ctx_lin) == 0, (
         "ReACT expects a fresh context."
     )

docs/tutorial.md

@@ -968,15 +968,15 @@ Let's look at how this agent is implemented in Mellea:
 ```python
 # file: https://github.com/generative-computing/mellea/blob/main/docs/examples/agents/react.py#L99
 def react(
-    m: mellea.MelleaSession,
-    goal: str,
-    react_toolbox: ReactToolbox,
-    budget: int=5,
+        m: mellea.MelleaSession,
+        goal: str,
+        react_toolbox: ReactToolbox,
+        budget: int = 5,
 ):
     assert m.ctx.is_chat_context, "ReACT requires a chat context."
-    test_ctx_lin = m.ctx.linearize()
+    test_ctx_lin = m.ctx.render_for_generation()
     assert (
-        test_ctx_lin is not None and len(test_ctx_lin) == 0
+            test_ctx_lin is not None and len(test_ctx_lin) == 0
     ), "ReACT expects a fresh context."
 
     # Construct the system prompt for ReACT.
@@ -1006,15 +1006,17 @@ def react(
             # model_options={mellea.backends.types.ModelOption.TOOLS: react_toolbox.tools_dict()},
             format=react_toolbox.tool_name_schema(),
         )
-        selected_tool: ReactTool = react_toolbox.get_tool_from_schema(act.content)
+        selected_tool: ReactTool = react_toolbox.get_tool_from_schema(
+            act.content)
         print(selected_tool.get_name())
 
         print(f"### Arguments for action")
         act_args = m.chat(
             "Choose arguments for the tool. Respond using JSON and include only the tool arguments in your response.",
             format=selected_tool.args_schema(),
         )
-        print(f"```json\n{json.dumps(json.loads(act_args.content), indent=2)}\n```")
+        print(
+            f"```json\n{json.dumps(json.loads(act_args.content), indent=2)}\n```")
 
         # TODO: handle exceptions.
         print("### Observation")

mellea/backends/formatter.py

@@ -166,7 +166,7 @@ def print_context(self, ctx: Context) -> str:
         )
         match ctx:
             case LinearContext():
-                linearized_ctx = ctx.linearize()
+                linearized_ctx = ctx.render_for_generation()
                 assert linearized_ctx is not None
                 return "".join([self.print(x) for x in linearized_ctx])
             case SimpleContext():

mellea/backends/huggingface.py

@@ -239,7 +239,7 @@ def _generate_from_context_alora(
                     "This code block should not execute unless there is a 'constraint' alora loaded."
                 )
         # Construct the linearized context. This is very similar to normal generation.
-        linearized_ctx = ctx.linearize()
+        linearized_ctx = ctx.render_for_generation()
         assert linearized_ctx is not None and len(linearized_ctx) > 1
         msgs = self.formatter.to_chat_messages(linearized_ctx)
         user_message, assistant_message = msgs[-2].content, msgs[-1].content
@@ -286,7 +286,7 @@ def _generate_from_context_standard(
         # Otherwise, we will linearize the context and treat it as a raw input.
         decoded_result: str | None = None
         if ctx.is_chat_context:
-            linearized_ctx = ctx.linearize()
+            linearized_ctx = ctx.render_for_generation()
             assert linearized_ctx is not None, (
                 "If ctx.is_chat_context, then the context should be linearizable."
             )

mellea/backends/ollama.py

@@ -263,7 +263,7 @@ def generate_from_chat_context(
         """
         model_opts = self._simplify_and_merge(model_options)
 
-        linearized_context = ctx.linearize()
+        linearized_context = ctx.render_for_generation()
         assert linearized_context is not None, (
             "Cannot generate from a non-linear context in a FormatterBackend."
         )

mellea/backends/openai.py

@@ -327,7 +327,7 @@ def _generate_from_chat_context_alora(
                 )
 
         # Construct the linearized context. This is very similar to normal generation.
-        linearized_ctx = ctx.linearize()
+        linearized_ctx = ctx.render_for_generation()
         assert linearized_ctx is not None and len(linearized_ctx) > 1
         msgs = self.formatter.to_chat_messages(linearized_ctx)
         user_message, assistant_message = msgs[-2].content, msgs[-1].content
@@ -362,7 +362,7 @@ def _generate_from_chat_context_standard(
         model_opts = self._simplify_and_merge(
             model_options, is_chat_context=ctx.is_chat_context
         )
-        linearized_context = ctx.linearize()
+        linearized_context = ctx.render_for_generation()
         assert linearized_context is not None, (
             "Cannot generate from a non-linear context in a FormatterBackend."
         )

mellea/backends/watsonx.py

@@ -220,7 +220,7 @@ def generate_from_chat_context(
             model_options, is_chat_context=ctx.is_chat_context
         )
 
-        linearized_context = ctx.linearize()
+        linearized_context = ctx.render_for_generation()
         assert linearized_context is not None, (
             "Cannot generate from a non-linear context in a FormatterBackend."
         )

mellea/stdlib/base.py

@@ -157,8 +157,13 @@ def _hash_for_kv_cache(self):
         ...
 
     @abc.abstractmethod
-    def linearize(self) -> list[Component | CBlock] | None:
-        """Provides a linear list of context components. This is not always possible, or None if that is not possible to construct."""
+    def render_for_generation(self) -> list[Component | CBlock] | None:
+        """Provides a linear list of context components to use for generation, or None if that is not possible to construct."""
+        ...
+
+    @abc.abstractmethod
+    def full_event_log(self) -> list[Component | CBlock]:
+        """Provides a list of all events stored in the context."""
         ...
 
     @abc.abstractmethod
@@ -262,6 +267,10 @@ def last_output_and_logs(
                 )
                 return last, log[0]
 
+    def full_event_log(self) -> list[Component | CBlock]:
+        """Returns the underlying _ctx."""
+        return self._ctx
+
     def last_turn(self):
         """The last input/output turn of the context."""
         if len(self._ctx) == 0:
@@ -335,8 +344,8 @@ def insert_turn(
         if turn.output:
             self.insert(turn.output, generate_logs=generate_logs)
 
-    def linearize(self) -> list[Component | CBlock] | None:
-        """Returns the underlying _ctx list."""
+    def render_for_generation(self) -> list[Component | CBlock] | None:
+        """Returns the underlying _ctx list for generation."""
         return self._ctx
 
     def is_chat_history(self):
@@ -372,7 +381,7 @@ def __init__(self):
         super().__init__()
         self.is_chat_context = True
 
-    def linearize(self) -> list[Component | CBlock] | None:
+    def render_for_generation(self) -> list[Component | CBlock] | None:
         """Uses _ctx ordering."""
         return []
 

mellea/stdlib/chat.py

@@ -111,10 +111,10 @@ def _to_msg(c: CBlock | Component | ModelOutputThunk) -> Message | None:
             case _:
                 return None
 
-    linearized_ctx = ctx.linearize()
-    if linearized_ctx is None:
+    all_ctx_events = ctx.full_event_log()
+    if all_ctx_events is None:
         raise Exception("Trying to cast a non-linear history into a chat history.")
     else:
-        history = [_to_msg(c) for c in linearized_ctx]
+        history = [_to_msg(c) for c in all_ctx_events]
         assert None not in history, "Could not render this context as a chat history."
         return history  # type: ignore