goose-py 0.7.3__tar.gz → 0.9.0__tar.gz

{goose_py-0.7.3 → goose_py-0.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: goose-py
-Version: 0.7.3
+Version: 0.9.0
 Summary: A tool for AI workflows based on human-computer collaboration and structured output.
 Author-email: Nash Taylor <nash@chelle.ai>, Joshua Cook <joshua@chelle.ai>, Michael Sankur <michael@chelle.ai>
 Requires-Python: >=3.12

{goose_py-0.7.3 → goose_py-0.9.0}/goose/__init__.py

@@ -1,4 +1,4 @@
 from goose._internal.agent import Agent
-from goose._internal.flow import flow
+from goose._internal.flow import flow, FlowArguments
 from goose._internal.result import Result, TextResult
 from goose._internal.task import task

{goose_py-0.7.3 → goose_py-0.9.0}/goose/_internal/flow.py

@@ -1,4 +1,4 @@
-from collections.abc import AsyncIterator, Awaitable, Callable
+from collections.abc import AsyncIterator, Callable
 from contextlib import asynccontextmanager
 from types import CodeType
 from typing import Protocol, overload
@@ -6,48 +6,62 @@ from typing import Protocol, overload
 from goose._internal.agent import Agent, IAgentLogger
 from goose._internal.conversation import Conversation
 from goose._internal.result import Result
-from goose._internal.state import FlowRun, get_current_flow_run, set_current_flow_run
+from goose._internal.state import FlowArguments, FlowRun, get_current_flow_run, set_current_flow_run
 from goose._internal.store import IFlowRunStore, InMemoryFlowRunStore
 from goose.errors import Honk
 
 
+class IGenerator[FlowArgumentsT: FlowArguments](Protocol):
+    __name__: str
+
+    async def __call__(self, *, flow_arguments: FlowArgumentsT, agent: Agent) -> None: ...
+
+
 class IAdapter[ResultT: Result](Protocol):
     __code__: CodeType
 
     async def __call__(self, *, conversation: Conversation[ResultT], agent: Agent) -> ResultT: ...
 
 
-class Flow[**P]:
+class Flow[FlowArgumentsT: FlowArguments]:
     def __init__(
         self,
-        fn: Callable[P, Awaitable[None]],
+        fn: IGenerator[FlowArgumentsT],
         /,
         *,
         name: str | None = None,
-        store: IFlowRunStore | None = None,
+        store: IFlowRunStore[FlowArgumentsT] | None = None,
         agent_logger: IAgentLogger | None = None,
     ) -> None:
         self._fn = fn
         self._name = name
         self._agent_logger = agent_logger
-        self._store = store or InMemoryFlowRunStore(flow_name=self.name)
+        self._store = store or InMemoryFlowRunStore(flow_name=self.name, flow_arguments_model=self.flow_arguments_model)
+
+    @property
+    def flow_arguments_model(self) -> type[FlowArgumentsT]:
+        arguments_model = self._fn.__annotations__.get("flow_arguments")
+        if arguments_model is None:
+            raise Honk("Flow function has an invalid signature. Must accept `flow_arguments` and `agent` as arguments.")
+
+        return arguments_model
 
     @property
     def name(self) -> str:
         return self._name or self._fn.__name__
 
     @property
-    def current_run(self) -> FlowRun:
+    def current_run(self) -> FlowRun[FlowArgumentsT]:
         run = get_current_flow_run()
         if run is None:
             raise Honk("No current flow run")
         return run
 
     @asynccontextmanager
-    async def start_run(self, *, run_id: str) -> AsyncIterator[FlowRun]:
+    async def start_run(self, *, run_id: str) -> AsyncIterator[FlowRun[FlowArgumentsT]]:
         existing_run = await self._store.get(run_id=run_id)
         if existing_run is None:
-            run = FlowRun()
+            run = FlowRun(flow_arguments_model=self.flow_arguments_model)
         else:
             run = existing_run
 
@@ -61,43 +75,42 @@ class Flow[**P]:
 
         set_current_flow_run(old_run)
 
-    async def generate(self, *args: P.args, **kwargs: P.kwargs) -> None:
+    async def generate(self, flow_arguments: FlowArgumentsT, /) -> None:
         flow_run = get_current_flow_run()
         if flow_run is None:
             raise Honk("No current flow run")
 
-        flow_run.set_flow_inputs(*args, **kwargs)
-        await self._fn(*args, **kwargs)
+        flow_run.set_flow_arguments(flow_arguments)
+        await self._fn(flow_arguments=flow_arguments, agent=flow_run.agent)
 
     async def regenerate(self) -> None:
         flow_run = get_current_flow_run()
         if flow_run is None:
             raise Honk("No current flow run")
 
-        flow_args, flow_kwargs = flow_run.flow_inputs
-        await self._fn(*flow_args, **flow_kwargs)
+        await self._fn(flow_arguments=flow_run.flow_arguments, agent=flow_run.agent)
 
 
 @overload
-def flow[**P](fn: Callable[P, Awaitable[None]], /) -> Flow[P]: ...
+def flow[FlowArgumentsT: FlowArguments](fn: IGenerator[FlowArgumentsT], /) -> Flow[FlowArgumentsT]: ...
 @overload
-def flow[**P](
+def flow[FlowArgumentsT: FlowArguments](
     *,
     name: str | None = None,
-    store: IFlowRunStore | None = None,
+    store: IFlowRunStore[FlowArgumentsT] | None = None,
     agent_logger: IAgentLogger | None = None,
-) -> Callable[[Callable[P, Awaitable[None]]], Flow[P]]: ...
-def flow[**P](
-    fn: Callable[P, Awaitable[None]] | None = None,
+) -> Callable[[IGenerator[FlowArgumentsT]], Flow[FlowArgumentsT]]: ...
+def flow[FlowArgumentsT: FlowArguments](
+    fn: IGenerator[FlowArgumentsT] | None = None,
     /,
     *,
     name: str | None = None,
-    store: IFlowRunStore | None = None,
+    store: IFlowRunStore[FlowArgumentsT] | None = None,
     agent_logger: IAgentLogger | None = None,
-) -> Flow[P] | Callable[[Callable[P, Awaitable[None]]], Flow[P]]:
+) -> Flow[FlowArgumentsT] | Callable[[IGenerator[FlowArgumentsT]], Flow[FlowArgumentsT]]:
     if fn is None:
 
-        def decorator(fn: Callable[P, Awaitable[None]]) -> Flow[P]:
+        def decorator(fn: IGenerator[FlowArgumentsT]) -> Flow[FlowArgumentsT]:
             return Flow(fn, name=name, store=store, agent_logger=agent_logger)
 
         return decorator

{goose_py-0.7.3 → goose_py-0.9.0}/goose/_internal/state.py

@@ -1,8 +1,8 @@
+import json
 from contextvars import ContextVar
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Self
+from typing import TYPE_CHECKING, Any, NewType, Self
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from goose._internal.agent import (
     Agent,
@@ -17,12 +17,11 @@ from goose.errors import Honk
 if TYPE_CHECKING:
     from goose._internal.task import Task
 
+SerializedFlowRun = NewType("SerializedFlowRun", str)
 
-@dataclass
-class FlowRunState:
-    node_states: dict[tuple[str, int], str]
-    flow_args: tuple[Any, ...]
-    flow_kwargs: dict[str, Any]
+
+class FlowArguments(BaseModel):
+    model_config = ConfigDict(frozen=True)
 
 
 class NodeState[ResultT: Result](BaseModel):
@@ -73,15 +72,15 @@ class NodeState[ResultT: Result](BaseModel):
         return self
 
 
-class FlowRun:
-    def __init__(self) -> None:
+class FlowRun[FlowArgumentsT: FlowArguments]:
+    def __init__(self, *, flow_arguments_model: type[FlowArgumentsT]) -> None:
         self._node_states: dict[tuple[str, int], str] = {}
         self._last_requested_indices: dict[str, int] = {}
         self._flow_name = ""
         self._id = ""
         self._agent: Agent | None = None
-        self._flow_args: tuple[Any, ...] | None = None
-        self._flow_kwargs: dict[str, Any] | None = None
+        self._flow_arguments: FlowArgumentsT | None = None
+        self._flow_arguments_model = flow_arguments_model
 
     @property
     def flow_name(self) -> str:
@@ -98,11 +97,11 @@ class FlowRun:
         return self._agent
 
     @property
-    def flow_inputs(self) -> tuple[tuple[Any, ...], dict[str, Any]]:
-        if self._flow_args is None or self._flow_kwargs is None:
+    def flow_arguments(self) -> FlowArgumentsT:
+        if self._flow_arguments is None:
             raise Honk("This Flow run has not been executed before")
 
-        return self._flow_args, self._flow_kwargs
+        return self._flow_arguments
 
     def get_all[R: Result](self, *, task: "Task[Any, R]") -> list[NodeState[R]]:
         matching_nodes: list[NodeState[R]] = []
@@ -122,9 +121,8 @@ class FlowRun:
                 last_hash=0,
             )
 
-    def set_flow_inputs(self, *args: Any, **kwargs: Any) -> None:
-        self._flow_args = args
-        self._flow_kwargs = kwargs
+    def set_flow_arguments(self, flow_arguments: FlowArgumentsT, /) -> None:
+        self._flow_arguments = flow_arguments
 
     def upsert_node_state(self, node_state: NodeState[Any], /) -> None:
         key = (node_state.task_name, node_state.index)
@@ -161,31 +159,35 @@ class FlowRun:
         if key in self._node_states:
             del self._node_states[key]
 
-    def dump(self) -> FlowRunState:
-        flow_args, flow_kwargs = self.flow_inputs
-
-        return FlowRunState(
-            node_states=self._node_states,
-            flow_args=flow_args,
-            flow_kwargs=flow_kwargs,
+    def dump(self) -> SerializedFlowRun:
+        formatted_node_states = {f"{k[0]},{k[1]}": v for k, v in self._node_states.items()}
+        return SerializedFlowRun(
+            json.dumps({"node_states": formatted_node_states, "flow_arguments": self.flow_arguments.model_dump()})
         )
 
     @classmethod
-    def load(cls, flow_run_state: FlowRunState, /) -> Self:
-        flow_run = cls()
-        flow_run._node_states = flow_run_state.node_states
-        flow_run._flow_args = flow_run_state.flow_args
-        flow_run._flow_kwargs = flow_run_state.flow_kwargs
+    def load(cls, *, serialized_flow_run: SerializedFlowRun, flow_arguments_model: type[FlowArgumentsT]) -> Self:
+        flow_run_state = json.loads(serialized_flow_run)
+        raw_node_states = flow_run_state["node_states"]
+        node_states: dict[tuple[str, int], str] = {}
+        for key, value in raw_node_states.items():
+            task_name, index = key.split(",")
+            node_states[(task_name, int(index))] = value
+        flow_arguments = flow_arguments_model.model_validate(flow_run_state["flow_arguments"])
+
+        flow_run = cls(flow_arguments_model=flow_arguments_model)
+        flow_run._node_states = node_states
+        flow_run._flow_arguments = flow_arguments
 
         return flow_run
 
 
-_current_flow_run: ContextVar[FlowRun | None] = ContextVar("current_flow_run", default=None)
+_current_flow_run: ContextVar[FlowRun[Any] | None] = ContextVar("current_flow_run", default=None)
 
 
-def get_current_flow_run() -> FlowRun | None:
+def get_current_flow_run() -> FlowRun[Any] | None:
     return _current_flow_run.get()
 
 
-def set_current_flow_run(flow_run: FlowRun | None) -> None:
+def set_current_flow_run(flow_run: FlowRun[Any] | None) -> None:
     _current_flow_run.set(flow_run)

goose_py-0.9.0/goose/_internal/store.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from typing import Protocol
+
+from goose._internal.flow import FlowRun
+from goose._internal.state import FlowArguments, SerializedFlowRun
+
+
+class IFlowRunStore[FlowArgumentsT: FlowArguments](Protocol):
+    def __init__(self, *, flow_name: str) -> None: ...
+    async def get(self, *, run_id: str) -> FlowRun[FlowArgumentsT] | None: ...
+    async def save(self, *, run: FlowRun[FlowArgumentsT]) -> None: ...
+    async def delete(self, *, run_id: str) -> None: ...
+
+
+class InMemoryFlowRunStore[FlowArgumentsT: FlowArguments](IFlowRunStore[FlowArgumentsT]):
+    def __init__(self, *, flow_name: str, flow_arguments_model: type[FlowArgumentsT]) -> None:
+        self._flow_name = flow_name
+        self._flow_arguments_model = flow_arguments_model
+        self._runs: dict[str, SerializedFlowRun] = {}
+
+    async def get(self, *, run_id: str) -> FlowRun[FlowArgumentsT] | None:
+        serialized_flow_run = self._runs.get(run_id)
+        if serialized_flow_run is not None:
+            return FlowRun.load(
+                serialized_flow_run=serialized_flow_run, flow_arguments_model=self._flow_arguments_model
+            )
+
+    async def save(self, *, run: FlowRun[FlowArgumentsT]) -> None:
+        self._runs[run.id] = run.dump()
+
+    async def delete(self, *, run_id: str) -> None:
+        self._runs.pop(run_id, None)

{goose_py-0.7.3 → goose_py-0.9.0}/goose/_internal/task.py

@@ -1,5 +1,5 @@
 from collections.abc import Awaitable, Callable
-from typing import overload
+from typing import Any, overload
 
 from goose._internal.agent import Agent, GeminiModel, SystemMessage, UserMessage
 from goose._internal.conversation import Conversation
@@ -43,7 +43,7 @@ class Task[**P, R: Result]:
         else:
             return state.result
 
-    async def jam(
+    async def refine(
         self,
         *,
         user_message: UserMessage,
@@ -111,7 +111,7 @@ class Task[**P, R: Result]:
         except TypeError:
             raise Honk(f"Unhashable argument to task {self.name}: {args} {kwargs}")
 
-    def __get_current_flow_run(self) -> FlowRun:
+    def __get_current_flow_run(self) -> FlowRun[Any]:
         run = get_current_flow_run()
         if run is None:
             raise Honk("No current flow run")

goose_py-0.9.0/goose/runs.py

@@ -0,0 +1,4 @@
+from goose._internal.state import FlowRun, SerializedFlowRun
+from goose._internal.store import IFlowRunStore
+
+__all__ = ["FlowRun", "IFlowRunStore", "SerializedFlowRun"]

{goose_py-0.7.3 → goose_py-0.9.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "goose-py"
-version = "0.7.3"
+version = "0.9.0"
 description = "A tool for AI workflows based on human-computer collaboration and structured output."
 readme = "README.md"
 authors = [
@@ -75,7 +75,8 @@ filterwarnings = [
     "ignore::SyntaxWarning",
     "ignore::UserWarning",
 ]
-addopts = "-v"
 testpaths = ["tests"]
 pythonpath = ["."]
 python_files = "test_*.py"
+asyncio_default_fixture_loop_scope = "session"
+asyncio_mode = "auto"

{goose_py-0.7.3 → goose_py-0.9.0}/tests/test_agent.py

@@ -3,9 +3,12 @@ from unittest.mock import Mock
 import pytest
 from pytest_mock import MockerFixture
 
-from goose import TextResult, flow, task
-from goose._internal.agent import Agent, AgentResponse, IAgentLogger
-from goose.agent import GeminiModel, TextMessagePart, UserMessage
+from goose import Agent, FlowArguments, TextResult, flow, task
+from goose.agent import AgentResponse, GeminiModel, IAgentLogger, TextMessagePart, UserMessage
+
+
+class TestFlowArguments(FlowArguments):
+    pass
 
 
 class MockLiteLLMResponse:
@@ -32,7 +35,7 @@ async def use_agent(*, agent: Agent) -> TextResult:
 
 
 @flow
-async def agent_flow(*, agent: Agent) -> None:
+async def agent_flow(*, flow_arguments: TestFlowArguments, agent: Agent) -> None:
     await use_agent(agent=agent)
 
 
@@ -44,7 +47,7 @@ class CustomLogger(IAgentLogger):
 
 
 @flow(agent_logger=CustomLogger())
-async def agent_flow_with_custom_logger(*, agent: Agent) -> None:
+async def agent_flow_with_custom_logger(*, flow_arguments: TestFlowArguments, agent: Agent) -> None:
     await use_agent(agent=agent)
 
 
@@ -52,7 +55,7 @@ async def agent_flow_with_custom_logger(*, agent: Agent) -> None:
 @pytest.mark.usefixtures("mock_litellm")
 async def test_agent() -> None:
     async with agent_flow.start_run(run_id="1") as run:
-        await agent_flow.generate(agent=run.agent)
+        await agent_flow.generate(TestFlowArguments())
 
     assert run.get(task=use_agent).result.text == "Hello"
 
@@ -60,7 +63,7 @@ async def test_agent() -> None:
 @pytest.mark.asyncio
 @pytest.mark.usefixtures("mock_litellm")
 async def test_agent_custom_logger() -> None:
-    async with agent_flow_with_custom_logger.start_run(run_id="1") as run:
-        await agent_flow_with_custom_logger.generate(agent=run.agent)
+    async with agent_flow_with_custom_logger.start_run(run_id="1"):
+        await agent_flow_with_custom_logger.generate(TestFlowArguments())
 
     assert len(CustomLogger.logged_responses) == 1

goose_py-0.9.0/tests/test_complex_flow_arguments.py

@@ -0,0 +1,21 @@
+import pytest
+
+from goose import Agent, FlowArguments, flow
+
+
+class MyFlowArguments(FlowArguments):
+    text: str
+
+
+@flow
+async def my_flow(*, flow_arguments: MyFlowArguments, agent: Agent) -> None:
+    pass
+
+
+@pytest.mark.asyncio
+async def test_my_flow() -> None:
+    async with my_flow.start_run(run_id="1"):
+        await my_flow.generate(MyFlowArguments(text="Hello"))
+
+    async with my_flow.start_run(run_id="1"):
+        await my_flow.generate(MyFlowArguments(text="Hello"))

{goose_py-0.7.3 → goose_py-0.9.0}/tests/test_downstream_task.py

@@ -3,7 +3,12 @@ import string
 
 import pytest
 
-from goose import Result, flow, task
+from goose import Agent, FlowArguments, Result, flow, task
+
+
+class MyFlowArguments(FlowArguments):
+    n_characters: int
+    n_duplicates: int
 
 
 class GeneratedWord(Result):
@@ -21,15 +26,15 @@ async def duplicate_word(*, word: str, times: int) -> GeneratedWord:
 
 
 @flow
-async def downstream_task() -> None:
-    word = await generate_random_word(n_characters=10)
-    await duplicate_word(word=word.word, times=10)
+async def downstream_task(*, flow_arguments: MyFlowArguments, agent: Agent) -> None:
+    word = await generate_random_word(n_characters=flow_arguments.n_characters)
+    await duplicate_word(word=word.word, times=flow_arguments.n_duplicates)
 
 
 @pytest.mark.asyncio
 async def test_downstream_task() -> None:
     async with downstream_task.start_run(run_id="1") as run:
-        await downstream_task.generate()
+        await downstream_task.generate(MyFlowArguments(n_characters=10, n_duplicates=10))
 
     duplicated_word = run.get(task=duplicate_word)
     assert len(duplicated_word.result.word) == 100

{goose_py-0.7.3 → goose_py-0.9.0}/tests/test_looping.py

@@ -1,6 +1,10 @@
 import pytest
 
-from goose import Result, flow, task
+from goose import Agent, FlowArguments, Result, flow, task
+
+
+class MyFlowArguments(FlowArguments):
+    pass
 
 
 class CourseObjective(Result):
@@ -41,7 +45,7 @@ async def quiz_question(*, outcome: str) -> QuizQuestion:
 
 
 @flow
-async def course_content() -> None:
+async def course_content(*, flow_arguments: MyFlowArguments, agent: Agent) -> None:
     objective = await course_objective()
     num_outcomes = await number_of_learning_outcomes()
     outcomes: list[LearningOutcome] = []
@@ -55,7 +59,7 @@ async def course_content() -> None:
 @pytest.mark.asyncio
 async def test_generate_course_content() -> None:
     async with course_content.start_run(run_id="1") as run:
-        await course_content.generate()
+        await course_content.generate(MyFlowArguments())
 
     quiz_questions = run.get_all(task=quiz_question)
     assert quiz_questions[0].result.question == "What is the meaning of Learn Python?"

goose_py-0.7.3/tests/test_jamming.py → goose_py-0.9.0/tests/test_refining.py

@@ -5,10 +5,14 @@ from unittest.mock import Mock
 import pytest
 from pytest_mock import MockerFixture
 
-from goose import Result, flow, task
+from goose import Agent, FlowArguments, Result, flow, task
 from goose.agent import SystemMessage, TextMessagePart, UserMessage
 
 
+class MyFlowArguments(FlowArguments):
+    pass
+
+
 class GeneratedWord(Result):
     word: str
 
@@ -37,7 +41,7 @@ async def make_sentence(*, words: list[GeneratedWord]) -> GeneratedSentence:
 
 
 @flow
-async def sentence() -> None:
+async def sentence(*, flow_arguments: MyFlowArguments, agent: Agent) -> None:
     words = [await generate_random_word(n_characters=10) for _ in range(3)]
     await make_sentence(words=words)
 
@@ -46,31 +50,30 @@ async def sentence() -> None:
 @pytest.mark.usefixtures("generate_random_word_adapter")
 async def test_jamming() -> None:
     async with sentence.start_run(run_id="1") as first_run:
-        await sentence.generate()
+        await sentence.generate(MyFlowArguments())
 
     initial_random_words = first_run.get_all(task=generate_random_word)
     assert len(initial_random_words) == 3
 
     # imagine this is a new process
     async with sentence.start_run(run_id="1") as second_run:
-        await generate_random_word.jam(
-            index=1,
+        await generate_random_word.refine(
             user_message=UserMessage(parts=[TextMessagePart(text="Change it")]),
             context=SystemMessage(parts=[TextMessagePart(text="Extra info")]),
         )
 
     random_words = second_run.get_all(task=generate_random_word)
     assert len(random_words) == 3
-    assert random_words[0].result.word != "__ADAPTED__"  # not adapted
-    assert random_words[1].result.word == "__ADAPTED__"  # adapted
+    assert random_words[0].result.word == "__ADAPTED__"  # adapted
+    assert random_words[1].result.word != "__ADAPTED__"  # not adapted
     assert random_words[2].result.word != "__ADAPTED__"  # not adapted
 
     # imagine this is a new process
     async with sentence.start_run(run_id="1") as third_run:
-        await sentence.generate()
+        await sentence.generate(MyFlowArguments())
 
     resulting_sentence = third_run.get(task=make_sentence)
     assert (
         resulting_sentence.result.sentence
-        == f"{initial_random_words[0].result.word} __ADAPTED__ {initial_random_words[2].result.word}"
+        == f"__ADAPTED__ {initial_random_words[1].result.word} {initial_random_words[2].result.word}"
     )

{goose_py-0.7.3 → goose_py-0.9.0}/tests/test_regenerate.py

@@ -3,10 +3,15 @@ import string
 
 import pytest
 
-from goose import Result, flow, task
+from goose import Agent, FlowArguments, Result, flow, task
 from goose.errors import Honk
 
 
+class MyFlowArguments(FlowArguments):
+    n_characters: int
+    times: int
+
+
 class GeneratedWord(Result):
     word: str
 
@@ -24,15 +29,15 @@ async def duplicate_word(*, word: str, times: int) -> GeneratedWord:
 
 
 @flow
-async def flow_with_arguments(*, n_characters: int, times: int) -> None:
-    word = await generate_random_word(n_characters=n_characters)
-    await duplicate_word(word=word.word, times=times)
+async def flow_with_arguments(*, flow_arguments: MyFlowArguments, agent: Agent) -> None:
+    word = await generate_random_word(n_characters=flow_arguments.n_characters)
+    await duplicate_word(word=word.word, times=flow_arguments.times)
 
 
 @pytest.mark.asyncio
 async def test_flow_arguments_in_run() -> None:
     async with flow_with_arguments.start_run(run_id="1") as run:
-        await flow_with_arguments.generate(n_characters=10, times=10)
+        await flow_with_arguments.generate(MyFlowArguments(n_characters=10, times=10))
 
     async with flow_with_arguments.start_run(run_id="1") as run:
         await flow_with_arguments.regenerate()

{goose_py-0.7.3 → goose_py-0.9.0}/tests/test_state.py

@@ -5,11 +5,15 @@ from unittest.mock import Mock
 import pytest
 from pytest_mock import MockerFixture
 
-from goose import Result, flow, task
+from goose import Agent, FlowArguments, Result, flow, task
 from goose._internal.types.agent import SystemMessage, TextMessagePart, UserMessage
 from goose.errors import Honk
 
 
+class MyFlowArguments(FlowArguments):
+    n_characters: int
+
+
 class GeneratedWord(Result):
     word: str
 
@@ -38,15 +42,15 @@ async def make_sentence(*, words: list[GeneratedWord]) -> GeneratedSentence:
 
 
 @flow
-async def with_state() -> None:
-    word = await generate_random_word(n_characters=10)
+async def with_state(*, flow_arguments: MyFlowArguments, agent: Agent) -> None:
+    word = await generate_random_word(n_characters=flow_arguments.n_characters)
     await make_sentence(words=[word])
 
 
 @pytest.mark.asyncio
 async def test_state_causes_caching() -> None:
     async with with_state.start_run(run_id="1") as run:
-        await with_state.generate()
+        await with_state.generate(MyFlowArguments(n_characters=10))
 
     random_word = run.get(task=generate_random_word).result.word
 
@@ -54,7 +58,7 @@ async def test_state_causes_caching() -> None:
         with_state.current_run
 
     async with with_state.start_run(run_id="1") as new_run:
-        await with_state.generate()
+        await with_state.generate(MyFlowArguments(n_characters=10))
 
     new_random_word = new_run.get(task=generate_random_word).result.word
 
@@ -65,10 +69,10 @@ async def test_state_causes_caching() -> None:
 @pytest.mark.usefixtures("generate_random_word_adapter")
 async def test_state_undo() -> None:
     async with with_state.start_run(run_id="2"):
-        await with_state.generate()
+        await with_state.generate(MyFlowArguments(n_characters=10))
 
     async with with_state.start_run(run_id="2"):
-        await generate_random_word.jam(
+        await generate_random_word.refine(
             index=0,
             user_message=UserMessage(parts=[TextMessagePart(text="Change it")]),
             context=SystemMessage(parts=[TextMessagePart(text="Extra info")]),
@@ -83,7 +87,7 @@ async def test_state_undo() -> None:
 @pytest.mark.asyncio
 async def test_state_edit() -> None:
     async with with_state.start_run(run_id="3"):
-        await with_state.generate()
+        await with_state.generate(MyFlowArguments(n_characters=10))
 
     async with with_state.start_run(run_id="3") as run:
         generate_random_word.edit(result=GeneratedWord(word="__EDITED__"), index=0)

{goose_py-0.7.3 → goose_py-0.9.0}/uv.lock

@@ -234,7 +234,7 @@ wheels = [
 
 [[package]]
 name = "goose-py"
-version = "0.7.3"
+version = "0.9.0"
 source = { editable = "." }
 dependencies = [
     { name = "jsonpath-ng" },

goose_py-0.7.3/goose/_internal/store.py

@@ -1,30 +0,0 @@
-from __future__ import annotations
-
-from typing import Protocol
-
-from goose._internal.flow import FlowRun
-from goose._internal.state import FlowRunState
-
-
-class IFlowRunStore(Protocol):
-    def __init__(self, *, flow_name: str) -> None: ...
-    async def get(self, *, run_id: str) -> FlowRun | None: ...
-    async def save(self, *, run: FlowRun) -> None: ...
-    async def delete(self, *, run_id: str) -> None: ...
-
-
-class InMemoryFlowRunStore(IFlowRunStore):
-    def __init__(self, *, flow_name: str) -> None:
-        self._flow_name = flow_name
-        self._runs: dict[str, FlowRunState] = {}
-
-    async def get(self, *, run_id: str) -> FlowRun | None:
-        state = self._runs.get(run_id)
-        if state is not None:
-            return FlowRun.load(state)
-
-    async def save(self, *, run: FlowRun) -> None:
-        self._runs[run.id] = run.dump()
-
-    async def delete(self, *, run_id: str) -> None:
-        self._runs.pop(run_id, None)

goose_py-0.7.3/goose/runs.py

@@ -1,4 +0,0 @@
-from goose._internal.state import FlowRun, FlowRunState
-from goose._internal.store import IFlowRunStore
-
-__all__ = ["FlowRun", "FlowRunState", "IFlowRunStore"]

goose_py-0.7.3/tests/conftest.py

@@ -1,10 +0,0 @@
-import asyncio
-
-import pytest
-
-
-@pytest.fixture(scope="session")
-def event_loop():
-    loop = asyncio.get_event_loop()
-    yield loop
-    loop.close()

goose_py-0.7.3/tests/test_complex_flow_arguments.py

@@ -1,22 +0,0 @@
-import pytest
-from pydantic import BaseModel
-
-from goose import Agent, flow
-
-
-class MyMessage(BaseModel):
-    text: str
-
-
-@flow
-async def my_flow(*, message: MyMessage, agent: Agent) -> None:
-    pass
-
-
-@pytest.mark.asyncio
-async def test_my_flow() -> None:
-    async with my_flow.start_run(run_id="1") as run:
-        await my_flow.generate(message=MyMessage(text="Hello"), agent=run.agent)
-
-    async with my_flow.start_run(run_id="1") as run:
-        await my_flow.generate(message=MyMessage(text="Hello"), agent=run.agent)