Merge branch 'main' into avi/fix-51-generate-logs

Author: nrfulton

docs/examples/generative_slots/generative_slots.py

@@ -16,20 +16,19 @@ def generate_summary(text: str) -> str:
 
 
 if __name__ == "__main__":
-    m = start_session()
-    sentiment_component = classify_sentiment(m, text="I love this!")
-    print("Output sentiment is : ", sentiment_component)
-
-    summary = generate_summary(
-        m,
-        text="""
-        The eagle rays are a group of cartilaginous fishes in the family Myliobatidae,
-        consisting mostly of large species living in the open ocean rather than on the sea bottom.
-        Eagle rays feed on mollusks, and crustaceans, crushing their shells with their flattened teeth.
-        They are excellent swimmers and are able to breach the water up to several meters above the
-        surface. Compared with other rays, they have long tails, and well-defined, rhomboidal bodies.
-        They are ovoviviparous, giving birth to up to six young at a time. They range from 0.48 to
-        5.1 m (1.6 to 16.7 ft) in length and 7 m (23 ft) in wingspan.
-        """,
-    )
-    print("Generated summary is :", summary)
+    with start_session():
+        sentiment_component = classify_sentiment(text="I love this!")
+        print("Output sentiment is : ", sentiment_component)
+
+        summary = generate_summary(
+            text="""
+            The eagle rays are a group of cartilaginous fishes in the family Myliobatidae,
+            consisting mostly of large species living in the open ocean rather than on the sea bottom.
+            Eagle rays feed on mollusks, and crustaceans, crushing their shells with their flattened teeth.
+            They are excellent swimmers and are able to breach the water up to several meters above the
+            surface. Compared with other rays, they have long tails, and well-defined, rhomboidal bodies.
+            They are ovoviviparous, giving birth to up to six young at a time. They range from 0.48 to
+            5.1 m (1.6 to 16.7 ft) in length and 7 m (23 ft) in wingspan.
+            """,
+        )
+        print("Generated summary is :", summary)

docs/examples/instruct_validate_repair/101_email.py

@@ -1,14 +1,17 @@
 # This is the 101 example for using `session` and `instruct`.
 # helper function to wrap text
 from docs.examples.helper import w
-from mellea import start_session
+from mellea import start_session, instruct
 from mellea.backends.types import ModelOption
 
 # 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})
+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.")
+    email_v1 = m.instruct("Write an email to invite all interns to the office party.")
 
 # print result
 print(f"***** email ****\n{w(email_v1)}\n*******")

docs/tutorial.md

@@ -309,6 +309,11 @@ final_options = {
 
 2. **Pushing and popping model state**. Sessions offer the ability to push and pop model state. This means you can temporarily change the `model_options` for a series of calls by pushing a new set of `model_options` and then revert those changes with a pop.
 
+#### System Messages
+In Mellea, `ModelOption.SYSTEM_PROMPT` is the recommended way to add/change the system message for a prompt. Setting it at the backend/session level will use the provided message as the system prompt for all future calls (just like any other model option). Similarly, you can specify the system prompt parameter for any session-level function (like `m.instruct`) to replace it for just that call. 
+
+Mellea recommends applying the system message this way because some model-provider apis don't properly serialize messages with the `system` role and expect them as a separate parameter.
+
 ### Conclusion
 
 We have now worked up from a simple "Hello, World" example to our first generative programming design pattern: **Instruct - Validate - Repair (IVR)**.

mellea/__init__.py

@@ -3,7 +3,7 @@
 import mellea.backends.model_ids as model_ids
 from mellea.stdlib.base import LinearContext, SimpleContext
 from mellea.stdlib.genslot import generative
-from mellea.stdlib.session import MelleaSession, start_session
+from mellea.stdlib.session import MelleaSession, start_session, instruct, chat, validate, query, transform
 
 __all__ = [
     "LinearContext",
@@ -12,4 +12,9 @@
     "generative",
     "model_ids",
     "start_session",
+    "instruct",
+    "chat",
+    "validate",
+    "query",
+    "transform"
 ]

mellea/stdlib/genslot.py

@@ -9,6 +9,7 @@
 from pydantic import BaseModel, Field, create_model
 
 from mellea.stdlib.base import Component, TemplateRepresentation
+from mellea.stdlib.session import get_session
 
 P = ParamSpec("P")
 R = TypeVar("R")
@@ -152,17 +153,19 @@ def __init__(self, func: Callable[P, R]):
         functools.update_wrapper(self, func)
 
     def __call__(
-        self, m, model_options: dict | None = None, *args: P.args, **kwargs: P.kwargs
+        self, m=None, model_options: dict | None = None, *args: P.args, **kwargs: P.kwargs
     ) -> R:
         """Call the generative slot.
 
         Args:
-            m: MelleaSession: A mellea session
+            m: MelleaSession: A mellea session (optional, uses context if None)
             **kwargs: Additional Kwargs to be passed to the func
 
         Returns:
             ModelOutputThunk: Output with generated Thunk.
         """
+        if m is None:
+            m = get_session()
         slot_copy = deepcopy(self)
         arguments = bind_function_arguments(self._function._func, *args, **kwargs)
         if arguments:

mellea/stdlib/session.py

@@ -3,7 +3,9 @@
 from __future__ import annotations
 
 from copy import deepcopy
-from typing import Any, Literal
+import contextvars
+from contextlib import contextmanager
+from typing import Any, Generator, Literal, Optional
 
 from mellea.backends import Backend, BaseModelSubclass
 from mellea.backends.formatter import FormatterBackend
@@ -33,6 +35,26 @@
 from mellea.stdlib.sampling import SamplingResult, SamplingStrategy
 
 
+# Global context variable for the context session
+_context_session: contextvars.ContextVar[Optional["MelleaSession"]] = contextvars.ContextVar(
+    "context_session", default=None
+)
+
+
+def get_session() -> "MelleaSession":
+    """Get the current session from context.
+
+    Raises:
+        RuntimeError: If no session is currently active.
+    """
+    session = _context_session.get()
+    if session is None:
+        raise RuntimeError(
+            "No active session found. Use 'with start_session(...):' to create one."
+        )
+    return session
+
+
 def backend_name_to_class(name: str) -> Any:
     """Resolves backend names to Backend classes."""
     if name == "ollama":
@@ -50,7 +72,6 @@ def backend_name_to_class(name: str) -> Any:
     else:
         return None
 
-
 def start_session(
     backend_name: Literal["ollama", "hf", "openai", "watsonx"] = "ollama",
     model_id: str | ModelIdentifier = IBM_GRANITE_3_3_8B,
@@ -59,14 +80,64 @@ def start_session(
     model_options: dict | None = None,
     **backend_kwargs,
 ) -> MelleaSession:
-    """Helper for starting a new mellea session.
+    """Start a new Mellea session. Can be used as a context manager or called directly.
+
+    This function creates and configures a new Mellea session with the specified backend
+    and model. When used as a context manager (with `with` statement), it automatically
+    sets the session as the current active session for use with convenience functions
+    like `instruct()`, `chat()`, `query()`, and `transform()`. When called directly,
+    it returns a session object that can be used directly.
 
     Args:
-        backend_name (str): ollama | hf | openai
-        model_id (ModelIdentifier): a `ModelIdentifier` from the mellea.backends.model_ids module
-        ctx (Optional[Context]): If not provided, a `LinearContext` is used.
-        model_options (Optional[dict]): Backend will be instantiated with these as its default, if provided.
-        backend_kwargs: kwargs that will be passed to the backend for instantiation.
+        backend_name: The backend to use. Options are:
+            - "ollama": Use Ollama backend for local models
+            - "hf" or "huggingface": Use HuggingFace transformers backend
+            - "openai": Use OpenAI API backend
+            - "watsonx": Use IBM WatsonX backend
+        model_id: Model identifier or name. Can be a `ModelIdentifier` from
+            mellea.backends.model_ids or a string model name.
+        ctx: Context manager for conversation history. Defaults to SimpleContext().
+            Use LinearContext() for chat-style conversations.
+        model_options: Additional model configuration options that will be passed
+            to the backend (e.g., temperature, max_tokens, etc.).
+        **backend_kwargs: Additional keyword arguments passed to the backend constructor.
+
+    Returns:
+        MelleaSession: A session object that can be used as a context manager
+        or called directly with session methods.
+
+    Usage:
+        # As a context manager (sets global session):
+        with start_session("ollama", "granite3.3:8b") as session:
+            result = instruct("Generate a story")  # Uses current session
+            # session is also available directly
+            other_result = session.chat("Hello")
+
+        # Direct usage (no global session set):
+        session = start_session("ollama", "granite3.3:8b")
+        result = session.instruct("Generate a story")
+        # Remember to call session.cleanup() when done
+        session.cleanup()
+
+    Examples:
+        # Basic usage with default settings
+        with start_session() as session:
+            response = instruct("Explain quantum computing")
+
+        # Using OpenAI with custom model options
+        with start_session("openai", "gpt-4", model_options={"temperature": 0.7}):
+            response = chat("Write a poem")
+
+        # Using HuggingFace with LinearContext for conversations
+        from mellea.stdlib.base import LinearContext
+        with start_session("hf", "microsoft/DialoGPT-medium", ctx=LinearContext()):
+            chat("Hello!")
+            chat("How are you?")  # Remembers previous message
+
+        # Direct usage without context manager
+        session = start_session()
+        response = session.instruct("Explain quantum computing")
+        session.cleanup()
     """
     backend_class = backend_name_to_class(backend_name)
     if backend_class is None:
@@ -77,7 +148,6 @@ def start_session(
     backend = backend_class(model_id, model_options=model_options, **backend_kwargs)
     return MelleaSession(backend, ctx)
 
-
 class MelleaSession:
     """Mellea sessions are a THIN wrapper around `m` convenience functions with NO special semantics.
 
@@ -104,6 +174,19 @@ def __init__(self, backend: Backend, ctx: Context | None = None):
         self.ctx = ctx if ctx is not None else SimpleContext()
         self._backend_stack: list[tuple[Backend, dict | None]] = []
         self._session_logger = FancyLogger.get_logger()
+        self._context_token = None
+
+    def __enter__(self):
+        """Enter context manager and set this session as the current global session."""
+        self._context_token = _context_session.set(self)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit context manager and cleanup session."""
+        self.cleanup()
+        if self._context_token is not None:
+            _context_session.reset(self._context_token)
+            self._context_token = None
 
     def _push_model_state(self, new_backend: Backend, new_model_opts: dict):
         """The backend and model options used within a `Context` can be temporarily changed. This method changes the model's backend and model_opts, while saving the current settings in the `self._backend_stack`.
@@ -134,6 +217,13 @@ def reset(self):
         """Reset the context state."""
         self.ctx.reset()
 
+    def cleanup(self) -> None:
+        """Clean up session resources."""
+        self.reset()
+        self._backend_stack.clear()
+        if hasattr(self.backend, "close"):
+            self.backend.close()
+
     def summarize(self) -> ModelOutputThunk:
         """Summarizes the current context."""
         raise NotImplementedError()
@@ -588,3 +678,29 @@ def last_prompt(self) -> str | list[dict] | None:
             if isinstance(last_el, GenerateLog):
                 prompt = last_el.prompt
         return prompt
+
+
+# Convenience functions that use the current session
+def instruct(description: str, **kwargs) -> ModelOutputThunk | SamplingResult:
+    """Instruct using the current session."""
+    return get_session().instruct(description, **kwargs)
+
+
+def chat(content: str, **kwargs) -> Message:
+    """Chat using the current session."""
+    return get_session().chat(content, **kwargs)
+
+
+def validate(reqs, **kwargs):
+    """Validate using the current session."""
+    return get_session().validate(reqs, **kwargs)
+
+
+def query(obj: Any, query_str: str, **kwargs) -> ModelOutputThunk:
+    """Query using the current session."""
+    return get_session().query(obj, query_str, **kwargs)
+
+
+def transform(obj: Any, transformation: str, **kwargs):
+    """Transform using the current session."""
+    return get_session().transform(obj, transformation, **kwargs)