diff --git a/docs/examples/generative_slots/generative_slots.py b/docs/examples/generative_slots/generative_slots.py index 24c113098..7e6d0a5e6 100644 --- a/docs/examples/generative_slots/generative_slots.py +++ b/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) diff --git a/docs/examples/instruct_validate_repair/101_email.py b/docs/examples/instruct_validate_repair/101_email.py index 03238bde8..62e9f12bc 100644 --- a/docs/examples/instruct_validate_repair/101_email.py +++ b/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*******") diff --git a/mellea/__init__.py b/mellea/__init__.py index cacbd5112..5cabfebdc 100644 --- a/mellea/__init__.py +++ b/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" ] diff --git a/mellea/stdlib/genslot.py b/mellea/stdlib/genslot.py index a7ba74091..4a48f44bd 100644 --- a/mellea/stdlib/genslot.py +++ b/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: diff --git a/mellea/stdlib/session.py b/mellea/stdlib/session.py index f5a6c59dd..e7751e6aa 100644 --- a/mellea/stdlib/session.py +++ b/mellea/stdlib/session.py @@ -2,7 +2,9 @@ from __future__ import annotations -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.aloras.huggingface.granite_aloras import add_granite_aloras @@ -35,6 +37,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": @@ -48,7 +70,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, @@ -57,14 +78,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: @@ -75,7 +146,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. @@ -102,6 +172,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`. @@ -132,6 +215,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() @@ -572,3 +662,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) diff --git a/test/stdlib_basics/test_contextual_session.py b/test/stdlib_basics/test_contextual_session.py new file mode 100644 index 000000000..976998311 --- /dev/null +++ b/test/stdlib_basics/test_contextual_session.py @@ -0,0 +1,215 @@ +import pytest +from typing import Literal +from mellea import generative, start_session, instruct, chat, validate, query, transform +from mellea.stdlib.base import ModelOutputThunk +from mellea.stdlib.session import get_session, MelleaSession +from mellea.stdlib.mify import mify, MifiedProtocol +from mellea.stdlib.requirement import req + + +@generative +def classify_sentiment(text: str) -> Literal["positive", "negative"]: ... + + +@generative +def generate_summary(text: str) -> str: ... + + +@mify(fields_include={"name", "age"}) +class TestPerson: + def __init__(self, name: str, age: int): + self.name = name + self.age = age + + def get_info(self) -> str: + """Get person information.""" + return f"{self.name} is {self.age} years old" + + +def test_basic_contextual_session(): + """Test basic contextual session usage with convenience functions.""" + with start_session(): + # Test instruct + result = instruct("Say hello") + assert isinstance(result, ModelOutputThunk) + assert result.value is not None + + # Test that we can get the session + current_session = get_session() + assert isinstance(current_session, MelleaSession) + + +def test_no_active_session_error(): + """Test error handling when no session is active.""" + with pytest.raises(RuntimeError, match="No active session found"): + get_session() + + with pytest.raises(RuntimeError, match="No active session found"): + instruct("test") + + with pytest.raises(RuntimeError, match="No active session found"): + chat("test") + + +def test_generative_with_contextual_session(): + """Test generative slots work with contextual sessions.""" + with start_session(): + # Test without explicit session parameter + result = classify_sentiment(text="I love this!") + assert result in ["positive", "negative"] + + # Test with summary generation + summary = generate_summary(text="A short text about something interesting.") + assert isinstance(summary, str) + assert len(summary) > 0 + + +def test_generative_backward_compatibility(): + """Test that generative slots still work with explicit session parameter.""" + with start_session() as m: + # Test old pattern still works + result = classify_sentiment(m, text="I love this!") + assert result in ["positive", "negative"] + + +def test_mify_with_contextual_session(): + """Test mify functionality with contextual sessions.""" + with start_session(): + person = TestPerson("Alice", 30) + assert isinstance(person, MifiedProtocol) + + # Test query functionality + query_result = query(person, "What is this person's name?") + assert isinstance(query_result, ModelOutputThunk) + + # Test transform functionality + transform_result = transform(person, "Make this person 5 years older") + # Transform can return either ModelOutputThunk or the tool output when tools are called + assert transform_result is not None + + +def test_nested_sessions(): + """Test nested sessions behavior.""" + with start_session() as outer_session: + outer_result = instruct("outer session test") + assert isinstance(outer_result, ModelOutputThunk) + + with start_session() as inner_session: + # Inner session should be active + current_session = get_session() + assert current_session is inner_session + + inner_result = instruct("inner session test") + assert isinstance(inner_result, ModelOutputThunk) + + # After inner session exits, outer should be active again + current_session = get_session() + assert current_session is outer_session + + +def test_session_cleanup(): + """Test session cleanup after context exit.""" + session_ref = None + with start_session() as m: + session_ref = m + instruct("test during session") + + # After exiting context, no session should be active + with pytest.raises(RuntimeError, match="No active session found"): + get_session() + + # Session should have been cleaned up + assert hasattr(session_ref, 'ctx') + + +def test_all_convenience_functions(): + """Test all convenience functions work within contextual session.""" + with start_session(): + # Test instruct + instruct_result = instruct("Generate a greeting") + assert isinstance(instruct_result, ModelOutputThunk) + + # Test chat + chat_result = chat("Hello there") + assert hasattr(chat_result, 'content') + + # Test validate + validation = validate([req("The response should be positive")]) + assert validation is not None + + # Test query with a mified object + test_person = TestPerson("Test", 42) + query_result = query(test_person, "What is the name?") + assert isinstance(query_result, ModelOutputThunk) + + # Test transform with a mified object + transform_result = transform(test_person, "Double the age") + assert transform_result is not None + + +def test_session_with_parameters(): + """Test contextual session with custom parameters.""" + with start_session(backend_name="ollama", model_id="granite3.3:8b") as m: + result = instruct("test with parameters") + assert isinstance(result, ModelOutputThunk) + assert isinstance(m, MelleaSession) + + +def test_multiple_sequential_sessions(): + """Test multiple sequential contextual sessions.""" + # First session + with start_session(): + result1 = instruct("first session") + assert isinstance(result1, ModelOutputThunk) + + # Ensure no session is active between contexts + with pytest.raises(RuntimeError, match="No active session found"): + get_session() + + # Second session + with start_session(): + result2 = instruct("second session") + assert isinstance(result2, ModelOutputThunk) + + +def test_contextual_session_with_mified_object_methods(): + """Test that mified objects work properly within contextual sessions.""" + with start_session(): + person = TestPerson("Bob", 25) + + # Test that mified object methods work + query_obj = person.get_query_object("What's the age?") + assert query_obj is not None + + transform_obj = person.get_transform_object("Make older") + assert transform_obj is not None + + # Test format_for_llm + llm_format = person.format_for_llm() + assert llm_format is not None + assert hasattr(llm_format, 'args') + + +def test_session_methods_with_mified_objects(): + """Test using session query/transform methods with mified objects.""" + with start_session() as m: + person = TestPerson("Charlie", 35) + + # Test session query method + query_result = m.query(person, "What is this person's age?") + assert isinstance(query_result, ModelOutputThunk) + + # Test session transform method + transform_result = m.transform(person, "Make this person younger") + # Transform can return either ModelOutputThunk or tool output when tools are called + assert transform_result is not None + + # Verify mified objects have query/transform object creation methods + assert hasattr(person, 'get_query_object') + assert hasattr(person, 'get_transform_object') + assert hasattr(person, '_query_type') + assert hasattr(person, '_transform_type') + + +if __name__ == "__main__": + pytest.main([__file__]) \ No newline at end of file