goose-py 0.8.0__py3-none-any.whl → 0.9.1__py3-none-any.whl

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/_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/_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/_internal/store.py

@@ -3,27 +3,30 @@ from __future__ import annotations
 from typing import Protocol
 
 from goose._internal.flow import FlowRun
-from goose._internal.state import FlowRunState
+from goose._internal.state import FlowArguments, SerializedFlowRun
 
 
-class IFlowRunStore(Protocol):
+class IFlowRunStore[FlowArgumentsT: FlowArguments](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 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(IFlowRunStore):
-    def __init__(self, *, flow_name: 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._runs: dict[str, FlowRunState] = {}
+        self._flow_arguments_model = flow_arguments_model
+        self._runs: dict[str, SerializedFlowRun] = {}
 
-    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 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) -> None:
+    async def save(self, *, run: FlowRun[FlowArgumentsT]) -> None:
         self._runs[run.id] = run.dump()
 
     async def delete(self, *, run_id: str) -> None:

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
@@ -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/flow.py

@@ -1,3 +1,3 @@
-from goose._internal.flow import Flow
+from goose._internal.flow import Flow, FlowArguments, IGenerator
 
-__all__ = ["Flow"]
+__all__ = ["Flow", "FlowArguments", "IGenerator"]

goose/runs.py

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

{goose_py-0.8.0.dist-info → goose_py-0.9.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: goose-py
-Version: 0.8.0
+Version: 0.9.1
 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.9.1.dist-info/RECORD

@@ -0,0 +1,18 @@
+goose/__init__.py,sha256=_qg6tofbArI5KO7fWqaFNY5IcbVELTd7oS6uYijPmzg,185
+goose/agent.py,sha256=g2tPFqEhqBABEjmpNJ2ShfjHDGzmeUXIgOZCKDZ2-40,600
+goose/errors.py,sha256=-0OyZQJWYTRw5YgnCB2_uorVaUsL6Z0QYQO2FqzCiyg,32
+goose/flow.py,sha256=mRyGvgyGgy39P5ZDc1L-8_XqPf8HCcrqGrTBDSswn4M,116
+goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+goose/runs.py,sha256=kAHw08ayGbwzSH_p-CfsQDvTcA4H09nUuq6zt3gIFxo,170
+goose/_internal/agent.py,sha256=l0pKfShovrs238sKAr-zubtcacYm82TGwQHcBWVJm2g,5875
+goose/_internal/conversation.py,sha256=1OZQ_N6QZE7L_ZpXG2bjoWkVQ-G7h0JvKkqswmQWG58,1202
+goose/_internal/flow.py,sha256=LX6fdDBBgGNPfpSIUFenJ8-EI6tnNRuGZfUpvSb1isg,4128
+goose/_internal/result.py,sha256=-eZJn-2sPo7rHZ38Sz6IAHXqiJ-Ss39esEoFGimJEBI,155
+goose/_internal/state.py,sha256=wJdXHGw_KPqc1Tuv8cHj1loqLct_KHWBNXzbjBv6GAo,6552
+goose/_internal/store.py,sha256=ZdQwYRe1PKkTUd58p0ja1KyyGhEYGVP4gWvYsqOEWuk,1356
+goose/_internal/task.py,sha256=vfmWfxhr76nSF7STMZISFqP9ugt3BkHzraMbEZKuozI,5196
+goose/_internal/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+goose/_internal/types/agent.py,sha256=rNVt2gEr_m4_8tGFgcdichpPp8xhOS5GY0kN2C4tiE8,2153
+goose_py-0.9.1.dist-info/METADATA,sha256=O2kYoXBBpkOUDX1ejc8dBUHv6GI86trC-QTb11HoSAs,441
+goose_py-0.9.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+goose_py-0.9.1.dist-info/RECORD,,

goose_py-0.8.0.dist-info/RECORD

@@ -1,18 +0,0 @@
-goose/__init__.py,sha256=mppYCowcZw9ke_4y1d1ayHwI3502LBaY959jdOVBPp0,170
-goose/agent.py,sha256=g2tPFqEhqBABEjmpNJ2ShfjHDGzmeUXIgOZCKDZ2-40,600
-goose/errors.py,sha256=-0OyZQJWYTRw5YgnCB2_uorVaUsL6Z0QYQO2FqzCiyg,32
-goose/flow.py,sha256=A1bzNIjnoVXRFm6LGhQglxVnKMP0vEVfvTubTol7Kww,58
-goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose/runs.py,sha256=HHcky_IbmY9HWBjpXJgOcH2Ko0N39qADsGIPR8QYpek,160
-goose/_internal/agent.py,sha256=l0pKfShovrs238sKAr-zubtcacYm82TGwQHcBWVJm2g,5875
-goose/_internal/conversation.py,sha256=1OZQ_N6QZE7L_ZpXG2bjoWkVQ-G7h0JvKkqswmQWG58,1202
-goose/_internal/flow.py,sha256=KGT6NpkMY8q_N1yKwWrxfTbhwcu5AwdHtgCPqCdL3F8,3266
-goose/_internal/result.py,sha256=-eZJn-2sPo7rHZ38Sz6IAHXqiJ-Ss39esEoFGimJEBI,155
-goose/_internal/state.py,sha256=mXwSkTQq5Y6xLSRgvfUPILmK5nbiA-sp8Nunjcra5n4,6045
-goose/_internal/store.py,sha256=vIxPIpechF_lEQlQ8JT1NDySDfHe3-eMHEWeTqVbscg,946
-goose/_internal/task.py,sha256=gKOheVs4xCvqDLazzgp5Ok6bLz-MuywViiEXWNOQ_gQ,5186
-goose/_internal/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose/_internal/types/agent.py,sha256=rNVt2gEr_m4_8tGFgcdichpPp8xhOS5GY0kN2C4tiE8,2153
-goose_py-0.8.0.dist-info/METADATA,sha256=bg599aFGIn1jHNInnQ36ie5bF-2sgQMEhvV2naYZqi8,441
-goose_py-0.8.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-goose_py-0.8.0.dist-info/RECORD,,