feat: new context, new sampling,.

docs/examples/aLora/101_example.py

@@ -1,6 +1,6 @@
 import time
 
-from mellea import LinearContext, MelleaSession
+from mellea import LegacyLinearContext, MelleaSession
 from mellea.backends.aloras.huggingface.granite_aloras import HFConstraintAlora
 from mellea.backends.cache import SimpleLRUCache
 from mellea.backends.huggingface import LocalHFBackend
@@ -22,7 +22,7 @@
 backend.add_alora(custom_stembolt_failure_constraint)
 
 # Create M session
-m = MelleaSession(backend, ctx=LinearContext())
+m = MelleaSession(backend, ctx=LegacyLinearContext())
 
 # define a requirement
 failure_check = ALoraRequirement(

docs/examples/agents/react.py

@@ -178,7 +178,7 @@ def react(
 
 
 if __name__ == "__main__":
-    m = mellea.start_session(ctx=mellea.stdlib.base.LinearContext())
+    m = mellea.start_session(ctx=mellea.stdlib.base.LegacyLinearContext())
 
     def zip_lookup_tool_fn(city: str):
         """Returns the ZIP code for the `city`."""

docs/examples/agents/react_instruct.py

@@ -187,7 +187,7 @@ def react(
 
 
 if __name__ == "__main__":
-    m = mellea.start_session(ctx=mellea.stdlib.base.LinearContext())
+    m = mellea.start_session(ctx=mellea.stdlib.base.LegacyLinearContext())
 
     def zip_lookup_tool_fn(city: str):
         """Returns the ZIP code for the `city`."""

docs/examples/generative_slots/generate_with_context.py

@@ -1,4 +1,4 @@
-from mellea import LinearContext, generative, start_session
+from mellea import LegacyLinearContext, generative, start_session
 from mellea.backends.types import ModelOption
 from mellea.stdlib.base import CBlock
 
@@ -34,7 +34,7 @@ def give_feedback(essay: str) -> list[str]:
 
 if __name__ == "__main__":
     m = start_session(
-        ctx=LinearContext(), model_options={ModelOption.MAX_NEW_TOKENS: 100}
+        ctx=LegacyLinearContext(), model_options={ModelOption.MAX_NEW_TOKENS: 100}
     )
 
     text = """

docs/examples/image_text_models/vision_ollama_chat.py

@@ -2,10 +2,10 @@
 
 from PIL import Image
 
-from mellea import LinearContext, start_session
+from mellea import LegacyLinearContext, start_session
 from mellea.stdlib.base import ImageBlock
 
-m = start_session(model_id="granite3.2-vision", ctx=LinearContext())
+m = start_session(model_id="granite3.2-vision", ctx=LegacyLinearContext())
 # m = start_session(model_id="llava", ctx=LinearContext())
 
 # load image

docs/examples/mify/rich_document_advanced.py

@@ -9,7 +9,7 @@
 from mellea.stdlib.base import ModelOutputThunk, TemplateRepresentation
 
 # Use a `SimpleContext` so that each LLM call is independent.
-m = mellea.start_session(backend_name="hf", ctx=mellea.SimpleContext())
+m = mellea.start_session(backend_name="hf", ctx=mellea.LegacySimpleContext())
 
 # 2. Let's import docling so that we can process pdf documents.
 

docs/examples/notebooks/context_example.ipynb

@@ -71,7 +71,7 @@
    "source": [
     "## Import Mellea and Start a Session with LinearContext\n",
     "\n",
-    "Up to this point we have used SimpleContext, a context manager that resets the chat message history on each model call. That is, the model's context is entirely determined by the current Component. \n",
+    "Up to this point we have used SimpleContext, a context manager that resets the chat message history on each model call. That is, the model's context is entirely determined by the current Component.\n",
     "\n",
     "Mellea also provides a LinearContext, which behaves like a chat history. We will use the LinearContext to interact with cat hmodels:"
    ]
@@ -84,9 +84,9 @@
    },
    "outputs": [],
    "source": [
-    "from mellea import LinearContext, start_session\n",
+    "from mellea import LegacyLinearContext, start_session\n",
     "\n",
-    "m = start_session(ctx=LinearContext())\n",
+    "m = start_session(ctx=LegacyLinearContext())\n",
     "m.chat(\"Make up a math problem.\")\n",
     "m.chat(\"Solve your math problem.\")\n",
     "print(m.ctx.last_output())\n",

docs/examples/notebooks/m_serve_example.ipynb

@@ -83,11 +83,11 @@
     "\n",
     "import mellea\n",
     "from cli.serve.models import ChatMessage\n",
-    "from mellea.stdlib.base import LinearContext, ModelOutputThunk\n",
+    "from mellea.stdlib.base import LegacyLinearContext, ModelOutputThunk\n",
     "from mellea.stdlib.requirement import Requirement, simple_validate\n",
     "from mellea.stdlib.sampling import RejectionSamplingStrategy, SamplingResult\n",
     "\n",
-    "session = mellea.start_session(ctx=LinearContext())\n",
+    "session = mellea.start_session(ctx=LegacyLinearContext())\n",
     "\n",
     "\n",
     "def validate_hi_bob(email: str) -> bool:\n",

docs/examples/safety.py/guardian.py

@@ -4,7 +4,12 @@
 from mellea.backends import model_ids
 from mellea.backends.dummy import DummyBackend
 from mellea.backends.ollama import OllamaModelBackend
-from mellea.stdlib.base import Context, ContextTurn, ModelOutputThunk, SimpleContext
+from mellea.stdlib.base import (
+    ContextTurn,
+    LegacyContext,
+    LegacySimpleContext,
+    ModelOutputThunk,
+)
 from mellea.stdlib.chat import Message
 from mellea.stdlib.safety.guardian import GuardianCheck, GuardianRisk
 

docs/examples/sessions/creating_a_new_type_of_session.py

@@ -3,7 +3,12 @@
 from mellea import MelleaSession
 from mellea.backends import Backend, BaseModelSubclass
 from mellea.backends.ollama import OllamaModelBackend
-from mellea.stdlib.base import CBlock, Context, LinearContext, ModelOutputThunk
+from mellea.stdlib.base import (
+    CBlock,
+    LegacyContext,
+    LegacyLinearContext,
+    ModelOutputThunk,
+)
 from mellea.stdlib.chat import Message
 from mellea.stdlib.requirement import Requirement, reqify
 from mellea.stdlib.safety.guardian import GuardianCheck, GuardianRisk
@@ -18,7 +23,7 @@ def __init__(
         self,
         requirements: list[str | Requirement],
         backend: Backend,
-        ctx: Context | None = None,
+        ctx: LegacyContext | None = None,
         check_immediately: bool = True,
     ):
         super().__init__(backend, ctx)
@@ -66,7 +71,7 @@ def chat(
 m = ChatCheckingSession(
     requirements=[GuardianCheck("jailbreak"), GuardianCheck("profanity")],
     backend=OllamaModelBackend(),
-    ctx=LinearContext(),
+    ctx=LegacyLinearContext(),
 )
 
 # You can run this code to see the immediate checks working.

docs/examples/tutorial/context_example.py

@@ -1,6 +1,6 @@
-from mellea import LinearContext, start_session
+from mellea import LegacyLinearContext, start_session
 
-m = start_session(ctx=LinearContext())
+m = start_session(ctx=LegacyLinearContext())
 m.chat("Make up a math problem.")
 m.chat("Solve your math problem.")
 

docs/examples/tutorial/m_serve_example.py

@@ -4,11 +4,11 @@
 
 import mellea
 from cli.serve.models import ChatMessage
-from mellea.stdlib.base import LinearContext, ModelOutputThunk
+from mellea.stdlib.base import LegacyLinearContext, ModelOutputThunk
 from mellea.stdlib.requirement import Requirement, simple_validate
 from mellea.stdlib.sampling import RejectionSamplingStrategy, SamplingResult
 
-session = mellea.start_session(ctx=LinearContext())
+session = mellea.start_session(ctx=LegacyLinearContext())
 
 
 def validate_hi_bob(email: str) -> bool:

mellea/__init__.py

@@ -1,7 +1,7 @@
 """Mellea is a library for building robust LLM applications."""
 
 import mellea.backends.model_ids as model_ids
-from mellea.stdlib.base import LinearContext, SimpleContext
+from mellea.stdlib.base import LegacyLinearContext, LegacySimpleContext
 from mellea.stdlib.genslot import generative
 from mellea.stdlib.session import (
     MelleaSession,
@@ -14,9 +14,9 @@
 )
 
 __all__ = [
-    "LinearContext",
+    "LegacyLinearContext",
+    "LegacySimpleContext",
     "MelleaSession",
-    "SimpleContext",
     "chat",
     "generative",
     "instruct",

mellea/backends/__init__.py

@@ -9,7 +9,13 @@
 
 from mellea.backends.model_ids import ModelIdentifier
 from mellea.backends.types import ModelOption
-from mellea.stdlib.base import CBlock, Component, Context, GenerateLog, ModelOutputThunk
+from mellea.stdlib.base import (
+    CBlock,
+    Component,
+    GenerateLog,
+    LegacyContext,
+    ModelOutputThunk,
+)
 
 BaseModelSubclass = TypeVar(
     "BaseModelSubclass", bound=pydantic.BaseModel
@@ -37,7 +43,7 @@ def __init__(
     def generate_from_context(
         self,
         action: Component | CBlock,
-        ctx: Context,
+        ctx: LegacyContext,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict | None = None,

mellea/backends/dummy.py

@@ -1,7 +1,13 @@
 """This module holds shim backends used for smoke tests."""
 
 from mellea.backends import Backend, BaseModelSubclass
-from mellea.stdlib.base import CBlock, Component, Context, GenerateLog, ModelOutputThunk
+from mellea.stdlib.base import (
+    CBlock,
+    Component,
+    GenerateLog,
+    LegacyContext,
+    ModelOutputThunk,
+)
 
 
 class DummyBackend(Backend):
@@ -19,7 +25,7 @@ def __init__(self, responses: list[str] | None):
     def generate_from_context(
         self,
         action: Component | CBlock,
-        ctx: Context,
+        ctx: LegacyContext,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict | None = None,

mellea/backends/formatter.py

@@ -17,10 +17,10 @@
 from mellea.stdlib.base import (
     CBlock,
     Component,
-    Context,
-    LinearContext,
+    LegacyContext,
+    LegacyLinearContext,
+    LegacySimpleContext,
     ModelOutputThunk,
-    SimpleContext,
     TemplateRepresentation,
 )
 from mellea.stdlib.chat import Message, ToolMessage
@@ -36,7 +36,7 @@ def print(self, c: Component | CBlock) -> str:
         ...
 
     @abc.abstractmethod
-    def print_context(self, ctx: Context) -> str:
+    def print_context(self, ctx: LegacyContext) -> str:
         """Renders a Context for input to a model."""
         ...
 
@@ -170,17 +170,17 @@ def _parse(
         else:
             return result
 
-    def print_context(self, ctx: Context) -> str:
+    def print_context(self, ctx: LegacyContext) -> str:
         """Renders a Context for input to a model."""
         assert not ctx.is_chat_context, (
             "Chat contexts should be handled in a backend by first using `Formatter.to_chat_messages` and then passing the dict to an API endpoint or using hf.apply_chat_template."
         )
         match ctx:
-            case LinearContext():
+            case LegacyLinearContext():
                 linearized_ctx = ctx.render_for_generation()
                 assert linearized_ctx is not None
                 return "".join([self.print(x) for x in linearized_ctx])
-            case SimpleContext():
+            case LegacySimpleContext():
                 raise Exception("Do not know how to handle a SimpleContext yet.")
             case _:
                 raise Exception(

mellea/backends/huggingface.py

@@ -47,9 +47,9 @@
 from mellea.stdlib.base import (
     CBlock,
     Component,
-    Context,
     GenerateLog,
     GenerateType,
+    LegacyContext,
     ModelOutputThunk,
     ModelToolCall,
 )
@@ -186,7 +186,7 @@ def alora_model(self, model: "aLoRAPeftModelForCausalLM | None"):  # noqa: UP037
     def generate_from_context(
         self,
         action: Component | CBlock,
-        ctx: Context,
+        ctx: LegacyContext,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict | None = None,
@@ -218,7 +218,7 @@ def generate_from_context(
     def _generate_from_context_alora(
         self,
         action: Component | CBlock,
-        ctx: Context,
+        ctx: LegacyContext,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict[str, Any],
@@ -265,7 +265,7 @@ def _generate_from_context_alora(
     def _generate_from_context_standard(
         self,
         action: Component | CBlock,
-        ctx: Context,
+        ctx: LegacyContext,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict[str, Any],

mellea/backends/litellm.py

@@ -31,9 +31,9 @@
 from mellea.stdlib.base import (
     CBlock,
     Component,
-    Context,
     GenerateLog,
     GenerateType,
+    LegacyContext,
     ModelOutputThunk,
     ModelToolCall,
 )
@@ -106,7 +106,7 @@ def __init__(
     def generate_from_context(
         self,
         action: Component | CBlock,
-        ctx: Context,
+        ctx: LegacyContext,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict | None = None,
@@ -208,7 +208,7 @@ def _make_backend_specific_and_remove(
     def _generate_from_chat_context_standard(
         self,
         action: Component | CBlock,
-        ctx: Context,
+        ctx: LegacyContext,
         *,
         format: type[BaseModelSubclass]
         | None = None,  # Type[BaseModelSubclass] is a class object of a subclass of BaseModel

mellea/backends/ollama.py

@@ -23,9 +23,9 @@
 from mellea.stdlib.base import (
     CBlock,
     Component,
-    Context,
     GenerateLog,
     GenerateType,
+    LegacyContext,
     ModelOutputThunk,
     ModelToolCall,
     TemplateRepresentation,
@@ -232,7 +232,7 @@ def _make_backend_specific_and_remove(
     def generate_from_context(
         self,
         action: Component | CBlock,
-        ctx: Context,
+        ctx: LegacyContext,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict | None = None,
@@ -253,7 +253,7 @@ def generate_from_context(
     def generate_from_chat_context(
         self,
         action: Component | CBlock,
-        ctx: Context,
+        ctx: LegacyContext,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict | None = None,