feat: new context, new sampling,. (#166)

* non-breaking:
Context -> LegacyContext

* non-breaking:
*Context -> Legacy*Context

* Context, SimpleContext, ChatContext (nee:LinearContext)

without log mechanics...

* expand -> add

* remove RootContext in favor of _is_root attribute.

* .

* actions_for_available_tools is back

* all backends adjusted to new context structure.

Session modifications left open.. (unsafe commit)

* feat: add top level funcs and start sample changes

* context adjustments and tests

* maybe breaking:

removing chat, instruct,.. from mellea/__init__.py

* untested session and sampling changes

* backends should add ACTION and MOT.

* as_chat_history adapted

* fix: ollama client instantiate to prevent stale event loops in httpx sessions

* fix: mypy issues

* Renaming:

ctx.as_list(last_n_components..)

ctx.render_for_generation() -->
ctx.view_for_generation()

* name cleanup

* adding contexts to sampling

* 1) handling sampling results and context

2) making rejection-sampling default

* adding warning

* name clarity

* remove unnecessary function

* updated tests:

- all: convert to new cotext
- contextual_session test is commented out with TODO
- session_ctx test is removed because logging changed.

* fix: watsonx errors

* fix: sampling and vision tests

* examples react - error to be looked into by @jal

* fixed react - thanks to @jal

* examples for alora (untested) and gen slots (tested)

* more examples updated

* more examples updated

* remove old context

* fix: best of n and misc changes

* fix: improve runtime of best of n sampling

* fix: add tests

* fix: remove references to LinearContext

---------

Co-authored-by: jakelorocco <[email protected]>

Author: HendrikStrobelt

docs/examples/aLora/101_example.py

@@ -1,10 +1,10 @@
 import time
 
-from mellea import LinearContext, MelleaSession
+from mellea import MelleaSession
 from mellea.backends.aloras.huggingface.granite_aloras import HFConstraintAlora
 from mellea.backends.cache import SimpleLRUCache
 from mellea.backends.huggingface import LocalHFBackend
-from mellea.stdlib.base import GenerateLog
+from mellea.stdlib.base import ChatContext, GenerateLog
 from mellea.stdlib.requirement import ALoraRequirement, Requirement
 
 # Define a backend and add the constraint aLora
@@ -22,7 +22,7 @@
 backend.add_alora(custom_stembolt_failure_constraint)
 
 # Create M session
-m = MelleaSession(backend, ctx=LinearContext())
+m = MelleaSession(backend, ctx=ChatContext())
 
 # define a requirement
 failure_check = ALoraRequirement(

docs/examples/agents/react.py

@@ -2,7 +2,7 @@
 import inspect
 import json
 from collections.abc import Callable
-from typing import Literal, Unpack
+from typing import Literal
 
 import pydantic
 from jinja2 import Template
@@ -13,6 +13,7 @@
 import mellea.stdlib
 import mellea.stdlib.base
 import mellea.stdlib.chat
+from mellea.stdlib.base import ChatContext
 
 react_system_template: Template = Template(
     """Answer the user's question as best you can.
@@ -83,7 +84,7 @@ def call_tool(self, tool: ReactTool, kwargs_json: str):
     def tool_name_schema(self):
         names = self.tool_names()
         fields = dict()
-        fields["tool"] = Literal[Unpack[names]]
+        fields["tool"] = Literal[*names]
         return pydantic.create_model("ToolSelectionSchema", **fields)
 
     def get_tool_from_schema(self, content: str):
@@ -103,7 +104,7 @@ def react(
     react_toolbox: ReactToolbox,
 ):
     assert m.ctx.is_chat_context, "ReACT requires a chat context."
-    test_ctx_lin = m.ctx.render_for_generation()
+    test_ctx_lin = m.ctx.view_for_generation()
     assert test_ctx_lin is not None and len(test_ctx_lin) == 0, (
         "ReACT expects a fresh context."
     )
@@ -114,8 +115,9 @@ def react(
     )
 
     # Add the system prompt and the goal to the chat history.
-    m.ctx.insert(mellea.stdlib.chat.Message(role="system", content=_sys_prompt))
-    m.ctx.insert(mellea.stdlib.chat.Message(role="user", content=f"{goal}"))
+    m.ctx = m.ctx.add(
+        mellea.stdlib.chat.Message(role="system", content=_sys_prompt)
+    ).add(mellea.stdlib.chat.Message(role="user", content=f"{goal}"))
 
     # The main ReACT loop as a dynamic program:
     # (  ?(not done) ;
@@ -156,7 +158,7 @@ def react(
 
         print("### Observation")
         tool_output = react_toolbox.call_tool(selected_tool, act_args.content)
-        m.ctx.insert(mellea.stdlib.chat.Message(role="tool", content=tool_output))
+        m.ctx = m.ctx.add(mellea.stdlib.chat.Message(role="tool", content=tool_output))
         print(tool_output)
 
         print("### Done Check")
@@ -178,7 +180,7 @@ def react(
 
 
 if __name__ == "__main__":
-    m = mellea.start_session(ctx=mellea.stdlib.base.LinearContext())
+    m = mellea.start_session(ctx=ChatContext())
 
     def zip_lookup_tool_fn(city: str):
         """Returns the ZIP code for the `city`."""

docs/examples/agents/react_instruct.py

@@ -2,7 +2,7 @@
 import inspect
 import json
 from collections.abc import Callable
-from typing import Literal, Unpack
+from typing import Literal
 
 import pydantic
 from jinja2 import Template
@@ -11,6 +11,7 @@
 import mellea.stdlib
 import mellea.stdlib.base
 import mellea.stdlib.chat
+from mellea.stdlib.base import ChatContext
 
 react_system_template: Template = Template(
     """Answer the user's question as best you can.
@@ -81,7 +82,7 @@ def call_tool(self, tool: ReactTool, kwargs_json: str):
     def tool_name_schema(self):
         names = self.tool_names()
         fields = dict()
-        fields["tool"] = Literal[Unpack[names]]
+        fields["tool"] = Literal[*names]
         return pydantic.create_model("ToolSelectionSchema", **fields)
 
     def get_tool_from_schema(self, content: str):
@@ -101,7 +102,7 @@ def react(
     react_toolbox: ReactToolbox,
 ):
     assert m.ctx.is_chat_context, "ReACT requires a chat context."
-    test_ctx_lin = m.ctx.render_for_generation()
+    test_ctx_lin = m.ctx.view_for_generation()
     assert test_ctx_lin is not None and len(test_ctx_lin) == 0, (
         "ReACT expects a fresh context."
     )
@@ -112,8 +113,9 @@ def react(
     )
 
     # Add the system prompt and the goal to the chat history.
-    m.ctx.insert(mellea.stdlib.chat.Message(role="system", content=_sys_prompt))
-    m.ctx.insert(mellea.stdlib.chat.Message(role="user", content=f"{goal}"))
+    m.ctx = m.ctx.add(
+        mellea.stdlib.chat.Message(role="system", content=_sys_prompt)
+    ).add(mellea.stdlib.chat.Message(role="user", content=f"{goal}"))
 
     # The main ReACT loop as a dynamic program:
     # (  ?(not done) ;
@@ -159,7 +161,7 @@ def react(
 
         print("### Observation")
         tool_output = react_toolbox.call_tool(selected_tool, act_args_val)
-        m.ctx.insert(mellea.stdlib.chat.Message(role="tool", content=tool_output))
+        m.ctx = m.ctx.add(mellea.stdlib.chat.Message(role="tool", content=tool_output))
         print(tool_output)
 
         print("### Done Check")
@@ -187,7 +189,7 @@ def react(
 
 
 if __name__ == "__main__":
-    m = mellea.start_session(ctx=mellea.stdlib.base.LinearContext())
+    m = mellea.start_session(ctx=ChatContext())
 
     def zip_lookup_tool_fn(city: str):
         """Returns the ZIP code for the `city`."""

docs/examples/generative_slots/generate_with_context.py

@@ -1,6 +1,6 @@
-from mellea import LinearContext, generative, start_session
+from mellea import generative, start_session
 from mellea.backends.types import ModelOption
-from mellea.stdlib.base import CBlock
+from mellea.stdlib.base import CBlock, ChatContext
 
 # Generative slots can be used with sessions that have context.
 # By utilizing context, you can change the results of several
@@ -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=ChatContext(), model_options={ModelOption.MAX_NEW_TOKENS: 100}
     )
 
     text = """
@@ -55,7 +55,7 @@ def give_feedback(essay: str) -> list[str]:
 
     # If you have a set of generative functions, you can tweak them all by
     # adding context to the session they are running in.
-    m.ctx.insert(
+    m.ctx = m.ctx.add(
         CBlock(
             "You are an elementary school teacher. "
             "Any grades and feedback that you give should keep that in mind. Remember to be "
@@ -74,7 +74,7 @@ def give_feedback(essay: str) -> list[str]:
 
     # And, let's reset the context and try a different grading style.
     m.reset()
-    m.ctx.insert(
+    m.ctx = m.ctx.add(
         CBlock(
             "You are a grammarian that is focused solely on spelling and syntax, "
             "not on the content of essays. When giving grades and feedback, focus "

docs/examples/helper/__init__.py

@@ -1 +1 @@
-from .helpers import Any, fill, w
+from .helpers import req_print, w

docs/examples/helper/helpers.py

@@ -1,7 +1,14 @@
 from textwrap import fill
 from typing import Any
 
+from mellea.stdlib.requirement import Requirement, ValidationResult
+
 
 # Just for printing stuff nicely...
 def w(x: Any) -> str:
     return fill(str(x), width=120, replace_whitespace=False)
+
+
+def req_print(rv_list: list[tuple[Requirement, ValidationResult]]) -> str:
+    parts = [f"{bool(rv[1])}\t: {rv[0].description}" for rv in rv_list]
+    return "\n".join(parts)

docs/examples/image_text_models/vision_ollama_chat.py

@@ -2,11 +2,11 @@
 
 from PIL import Image
 
-from mellea import LinearContext, start_session
-from mellea.stdlib.base import ImageBlock
+from mellea import start_session
+from mellea.stdlib.base import ChatContext, ImageBlock
 
-m = start_session(model_id="granite3.2-vision", ctx=LinearContext())
-# m = start_session(model_id="llava", ctx=LinearContext())
+m = start_session(model_id="granite3.2-vision", ctx=ChatContext())
+# m = start_session(model_id="llava", ctx=ChatContext())
 
 # load image
 test_img = Image.open("pointing_up.jpg")

docs/examples/instruct_validate_repair/101_email.py

@@ -1,30 +1,19 @@
 # This is the 101 example for using `session` and `instruct`.
 # helper function to wrap text
 from docs.examples.helper import w
-from mellea import instruct, start_session
+from mellea import start_session
 from mellea.backends.types import ModelOption
 
-# create a session using Granite 3.3 8B on Ollama and a simple context [see below]
-with start_session(model_options={ModelOption.MAX_NEW_TOKENS: 200}):
-    # write an email
-    email_v1 = instruct("Write an email to invite all interns to the office party.")
-
 with start_session(model_options={ModelOption.MAX_NEW_TOKENS: 200}) as m:
     # write an email
     email_v1 = m.instruct("Write an email to invite all interns to the office party.")
+    print(m.last_prompt())
 
 # print result
 print(f"***** email ****\n{w(email_v1)}\n*******")
 
 # ************** END *************
 
-
-# # optionally: print the debug log for the last instruction on the context
-# from mellea.stdlib.base import GenerateLog
-# _, log = m.ctx.last_output_and_logs()
-# if isinstance(log, GenerateLog): # should be
-#     print(f"Prompt:\n{w(log.prompt)}") # print prompt
-
 # # start_session() is equivalent to:
 # from mellea.backends import model_ids
 # from mellea.backends.ollama import OllamaModelBackend

docs/examples/instruct_validate_repair/101_email_with_requirements.py

@@ -5,7 +5,7 @@
 # create a session using Granite 3.3 8B on Ollama and a simple context [see below]
 m = start_session(model_options={ModelOption.MAX_NEW_TOKENS: 200})
 
-# write an email
+# write an email with automatic requirement checking.
 email_v1 = m.instruct(
     "Write an email to invite all interns to the office party.",
     requirements=["be formal", "Use 'Dear interns' as greeting."],

docs/examples/instruct_validate_repair/101_email_with_validate.py

@@ -1,19 +1,29 @@
-from docs.examples.helper import w
+from docs.examples.helper import req_print, w
 from mellea import start_session
 from mellea.backends.types import ModelOption
 from mellea.stdlib.sampling import RejectionSamplingStrategy
 
 # create a session using Granite 3.3 8B on Ollama and a simple context [see below]
 m = start_session(model_options={ModelOption.MAX_NEW_TOKENS: 200})
 
-email_v1 = m.instruct(
+email_v2_samples = m.instruct(
     "Write an email to invite all interns to the office party.",
     requirements=["be formal", "Use 'Dear interns' as greeting."],
     strategy=RejectionSamplingStrategy(loop_budget=3),
+    return_sampling_results=True,
 )
 
-# print result
-print(f"***** email ****\n{w(email_v1)}\n*******")
+if email_v2_samples.success:
+    print(f"Success: \n{w(email_v2_samples.result)}")
+    print(
+        f"===> Requirement for this sample: \n{req_print(email_v2_samples.sample_validations[-1])}"
+    )
+else:
+    print(f"Failure: \n{w(email_v2_samples.result)}")
+    selected_index = email_v2_samples.sample_generations.index(email_v2_samples.result)
+    print(
+        f"===> Requirement for this sample: \n{req_print(email_v2_samples.sample_validations[selected_index])}"
+    )
 
 # # [optional] get logs for all loops:
 # from mellea.stdlib.base import GenerateLog