ollama generate walk.

TODO-nrf: we need to add generate walks to every generation call.

Author: nrfulton

docs/examples/melp/lazy.py

@@ -7,23 +7,28 @@
 backend = OllamaModelBackend("granite4:latest")
 
 
-async def main(backend: Backend, ctx: Context):
-    s1 = CBlock("What is 1+1? Respond with the number only.")
-    s1_out, _ = await backend.generate_from_context(action=s1, ctx=SimpleContext())
-
-    s2 = CBlock("What is 2+2? Respond with the number only.")
-    s2_out, _ = await backend.generate_from_context(action=s2, ctx=SimpleContext())
-
-    sc1 = SimpleComponent(
-        instruction="What is x+y? Respond with the number only", x=s1_out, y=s2_out
+async def fib(backend: Backend, ctx: Context, x: CBlock, y: CBlock) -> ModelOutputThunk:
+    sc = SimpleComponent(
+        instruction="What is x+y? Respond with the number only.", x=x, y=y
     )
+    mot, _ = await backend.generate_from_context(action=sc, ctx=SimpleContext())
+    return mot
 
-    print(await s1_out.avalue())
-    print(await s2_out.avalue())
 
-    sc1_out, _ = await backend.generate_from_context(action=sc1, ctx=SimpleContext())
-
-    print(await sc1_out.avalue())
+async def main(backend: Backend, ctx: Context):
+    fibs = []
+    for i in range(100):
+        if i == 0 or i == 1:
+            fibs.append(CBlock(f"{i + 1}"))
+        else:
+            fibs.append(await fib(backend, ctx, fibs[i - 1], fibs[i - 2]))
+
+    for x in fibs:
+        match x:
+            case ModelOutputThunk():
+                print(await x.avalue())
+            case CBlock():
+                print(x.value)
 
 
 asyncio.run(main(backend, SimpleContext()))

docs/examples/melp/lazy_fib.py

@@ -0,0 +1,38 @@
+import asyncio
+from mellea.stdlib.span import Span, SimpleComponent
+from mellea.stdlib.base import SimpleContext, Context, CBlock, ModelOutputThunk
+from mellea.stdlib.requirement import Requirement
+from mellea.backends import Backend
+from mellea.backends.ollama import OllamaModelBackend
+from typing import Tuple
+
+backend = OllamaModelBackend("granite4:latest")
+
+
+async def fib(backend: Backend, ctx: Context, x: CBlock, y: CBlock) -> ModelOutputThunk:
+    sc = SimpleComponent(
+        instruction="What is x+y? Respond with the number only.", x=x, y=y
+    )
+    mot, _ = await backend.generate_from_context(action=sc, ctx=SimpleContext())
+    return mot
+
+
+async def fib_main(backend: Backend, ctx: Context):
+    fibs = []
+    for i in range(20):
+        if i == 0 or i == 1:
+            fibs.append(CBlock(f"{i}"))
+        else:
+            mot = await fib(backend, ctx, fibs[i - 1], fibs[i - 2])
+            fibs.append(mot)
+
+    for x in enumerate(fibs):
+        match x:
+            case ModelOutputThunk():
+                n = await x.avalue()
+                print(n)
+            case CBlock():
+                print(x.value)
+
+
+asyncio.run(fib_main(backend, SimpleContext()))

docs/examples/melp/lazy_fib_sample.py

@@ -0,0 +1,61 @@
+import asyncio
+from mellea.stdlib.span import Span, SimpleComponent
+from mellea.stdlib.base import SimpleContext, Context, CBlock, ModelOutputThunk
+from mellea.stdlib.requirement import Requirement
+from mellea.backends import Backend
+from mellea.backends.ollama import OllamaModelBackend
+from typing import Tuple
+
+backend = OllamaModelBackend("granite4:latest")
+
+
+async def _fib_sample(
+    backend: Backend, ctx: Context, x: CBlock, y: CBlock
+) -> ModelOutputThunk | None:
+    sc = SimpleComponent(
+        instruction="What is x+y? Respond with the number only.", x=x, y=y
+    )
+    answer_mot, _ = await backend.generate_from_context(action=sc, ctx=SimpleContext())
+
+    # This is a fundamental thing: it means computation must occur.
+    # We need to be able to read this off at c.g. construction time.
+    value = await answer_mot.avalue()
+
+    try:
+        int(value)
+        return answer_mot
+    except:
+        return None
+
+
+async def fib_sampling_version(
+    backend: Backend, ctx: Context, x: CBlock, y: CBlock
+) -> ModelOutputThunk | None:
+    for i in range(5):
+        sample = await _fib_sample(backend, ctx, x, y)
+        if sample is not None:
+            return sample
+        else:
+            continue
+    return None
+
+
+async def fib_sampling_version_main(backend: Backend, ctx: Context):
+    fibs = []
+    for i in range(20):
+        if i == 0 or i == 1:
+            fibs.append(CBlock(f"{i}"))
+        else:
+            mot = await fib_sampling_version(backend, ctx, fibs[i - 1], fibs[i - 2])
+            fibs.append(mot)
+
+    for x_i, x in enumerate(fibs):
+        match x:
+            case ModelOutputThunk():
+                n = await x.avalue()
+                print(n)
+            case CBlock():
+                print(x.value)
+
+
+asyncio.run(fib_sampling_version_main(backend, SimpleContext()))

docs/examples/melp/states.py

@@ -0,0 +1,46 @@
+import mellea
+from mellea.stdlib.base import CBlock, Context, SimpleContext
+from mellea.stdlib.span import Span, SimpleComponent
+from mellea.backends import Backend
+from mellea.backends.ollama import OllamaModelBackend
+import asyncio
+
+
+async def main(backend: Backend, ctx: Context):
+    a_states = "Alaska,Arizona,Arkansas".split(",")
+    m_states = "Missouri", "Minnesota", "Montana", "Massachusetts"
+
+    a_state_pops = dict()
+    for state in a_states:
+        a_state_pops[state], _ = await backend.generate_from_context(
+            CBlock(f"What is the population of {state}? Respond with an integer only."),
+            SimpleContext(),
+        )
+    a_total_pop = SimpleComponent(
+        instruction=CBlock(
+            "What is the total population of these states? Respond with an integer only."
+        ),
+        **a_state_pops,
+    )
+    a_state_total, _ = await backend.generate_from_context(a_total_pop, SimpleContext())
+
+    m_state_pops = dict()
+    for state in m_states:
+        m_state_pops[state], _ = await backend.generate_from_context(
+            CBlock(f"What is the population of {state}? Respond with an integer only."),
+            SimpleContext(),
+        )
+    m_total_pop = SimpleComponent(
+        instruction=CBlock(
+            "What is the total population of these states? Respond with an integer only."
+        ),
+        **m_state_pops,
+    )
+    m_state_total, _ = await backend.generate_from_context(m_total_pop, SimpleContext())
+
+    print(await a_state_total.avalue())
+    print(await m_state_total.avalue())
+
+
+backend = OllamaModelBackend(model_id="granite4:latest")
+asyncio.run(main(backend, SimpleContext()))

mellea/backends/_utils.py

@@ -1,13 +1,20 @@
 from __future__ import annotations
 
 import inspect
+import itertools
 from collections.abc import Callable
 from typing import Any, Literal
 
 from mellea.backends.formatter import Formatter
 from mellea.backends.tools import parse_tools
 from mellea.helpers.fancy_logger import FancyLogger
-from mellea.stdlib.base import CBlock, Component, Context, ModelToolCall
+from mellea.stdlib.base import (
+    CBlock,
+    Component,
+    Context,
+    ModelOutputThunk,
+    ModelToolCall,
+)
 from mellea.stdlib.chat import Message
 from mellea.stdlib.requirement import ALoraRequirement, LLMaJRequirement, Requirement
 
@@ -80,3 +87,15 @@ def to_tool_calls(
     if len(model_tool_calls) > 0:
         return model_tool_calls
     return None
+
+
+def generate_walk(c: CBlock | Component | ModelOutputThunk) -> list[ModelOutputThunk]:
+    """Returns the generation walk ordering for a Span."""
+    match c:
+        case ModelOutputThunk() if not c.is_computed():
+            return [c]
+        case CBlock():
+            return []
+        case Component():
+            parts_walk = [generate_walk(p) for p in c.parts()]
+            return itertools.chain.from_iterable(parts_walk)  # aka flatten

mellea/backends/ollama.py

@@ -11,6 +11,7 @@
 
 import mellea.backends.model_ids as model_ids
 from mellea.backends import BaseModelSubclass
+from mellea.backends._utils import generate_walk
 from mellea.backends.formatter import Formatter, FormatterBackend, TemplateFormatter
 from mellea.backends.model_ids import ModelIdentifier
 from mellea.backends.tools import (
@@ -294,6 +295,12 @@ async def generate_from_chat_context(
         Raises:
             RuntimeError: If not called from a thread with a running event loop.
         """
+        # Start by awaiting any necessary computation.
+        _computed = [await todo.avalue() for todo in generate_walk(action)]
+        FancyLogger.get_logger().info(
+            f"generate_from_chat_context awaited on {len(_computed)} uncomputed mots."
+        )
+
         model_opts = self._simplify_and_merge(model_options)
 
         linearized_context = ctx.view_for_generation()
@@ -408,9 +415,14 @@ async def generate_from_raw(
 
         model_opts = self._simplify_and_merge(model_options)
 
+        for act in actions:
+            for todo in generate_walk(act):
+                await todo.avalue()
+
+        prompts = [self.formatter.print(action) for action in actions]
+
         # Ollama doesn't support "batching". There's some ability for concurrency. Use that here.
         # See https://github.com/ollama/ollama/blob/main/docs/faq.md#how-does-ollama-handle-concurrent-requests.
-        prompts = [self.formatter.print(action) for action in actions]
 
         # Run async so that we can make use of Ollama's concurrency.
         coroutines: list[Coroutine[Any, Any, ollama.GenerateResponse]] = []

mellea/stdlib/span/__init__.py

@@ -23,14 +23,14 @@ def __init__(self, **kwargs):
 
     def parts(self):
         """Returns the values of the kwargs."""
-        return self._kwargs.values()
+        return list(self._kwargs.values())
 
     def _kwargs_type_check(self, kwargs):
         for key in kwargs.keys():
             value = kwargs[key]
             assert issubclass(type(value), Component) or issubclass(
                 type(value), CBlock
-            ), f"Expected span but found {type(value)}"
+            ), f"Expected span but found {type(value)} of value: {value}"
             assert type(key) is str
         return True
 
@@ -41,9 +41,23 @@ def make_simple_string(kwargs):
             [f"<|{key}|>{value}</|{key}|>" for (key, value) in kwargs.items()]
         )
 
+    @staticmethod
+    def make_json_string(kwargs):
+        """Uses json."""
+        str_args = dict()
+        for key in kwargs.keys():
+            match kwargs[key]:
+                case ModelOutputThunk() | CBlock():
+                    str_args[key] = kwargs[key].value
+                case Component():
+                    str_args[key] = kwargs[key].format_for_llm()
+        import json
+
+        return json.dumps(str_args)
+
     def format_for_llm(self):
         """Uses a string rep."""
-        return SimpleComponent.make_simple_string(self._kwargs)
+        return SimpleComponent.make_json_string(self._kwargs)
         # """ Uses a simple tagging structure that needs to be changed in the future. """
         # return TemplateRepresentation(
         #     obj=self,