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

goose/_internal/flow.py

@@ -30,13 +30,13 @@ class Flow[FlowArgumentsT: FlowArguments]:
         /,
         *,
         name: str | None = None,
-        store: IFlowRunStore[FlowArgumentsT] | None = None,
+        store: IFlowRunStore | 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, flow_arguments_model=self.flow_arguments_model)
+        self._store = store or InMemoryFlowRunStore(flow_name=self.name)
 
     @property
     def flow_arguments_model(self) -> type[FlowArgumentsT]:
@@ -59,18 +59,20 @@ class Flow[FlowArgumentsT: FlowArguments]:
 
     @asynccontextmanager
     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(flow_arguments_model=self.flow_arguments_model)
+        existing_serialized_run = await self._store.get(run_id=run_id)
+        if existing_serialized_run is not None:
+            run = FlowRun.load(
+                serialized_flow_run=existing_serialized_run, flow_arguments_model=self.flow_arguments_model
+            )
         else:
-            run = existing_run
+            run = FlowRun(flow_arguments_model=self.flow_arguments_model)
 
         old_run = get_current_flow_run()
         set_current_flow_run(run)
 
         run.start(flow_name=self.name, run_id=run_id, agent_logger=self._agent_logger)
         yield run
-        await self._store.save(run=run)
+        await self._store.save(run_id=run_id, run=run.dump())
         run.end()
 
         set_current_flow_run(old_run)
@@ -97,7 +99,7 @@ def flow[FlowArgumentsT: FlowArguments](fn: IGenerator[FlowArgumentsT], /) -> Fl
 def flow[FlowArgumentsT: FlowArguments](
     *,
     name: str | None = None,
-    store: IFlowRunStore[FlowArgumentsT] | None = None,
+    store: IFlowRunStore | None = None,
     agent_logger: IAgentLogger | None = None,
 ) -> Callable[[IGenerator[FlowArgumentsT]], Flow[FlowArgumentsT]]: ...
 def flow[FlowArgumentsT: FlowArguments](
@@ -105,7 +107,7 @@ def flow[FlowArgumentsT: FlowArguments](
     /,
     *,
     name: str | None = None,
-    store: IFlowRunStore[FlowArgumentsT] | None = None,
+    store: IFlowRunStore | None = None,
     agent_logger: IAgentLogger | None = None,
 ) -> Flow[FlowArgumentsT] | Callable[[IGenerator[FlowArgumentsT]], Flow[FlowArgumentsT]]:
     if fn is None:

goose/_internal/state.py

@@ -166,7 +166,9 @@ class FlowRun[FlowArgumentsT: FlowArguments]:
         )
 
     @classmethod
-    def load(cls, *, serialized_flow_run: SerializedFlowRun, flow_arguments_model: type[FlowArgumentsT]) -> Self:
+    def load[T: FlowArguments](
+        cls, *, serialized_flow_run: SerializedFlowRun, flow_arguments_model: type[T]
+    ) -> "FlowRun[T]":
         flow_run_state = json.loads(serialized_flow_run)
         raw_node_states = flow_run_state["node_states"]
         node_states: dict[tuple[str, int], str] = {}
@@ -175,7 +177,7 @@ class FlowRun[FlowArgumentsT: FlowArguments]:
             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 = FlowRun(flow_arguments_model=flow_arguments_model)
         flow_run._node_states = node_states
         flow_run._flow_arguments = flow_arguments
 

goose/_internal/store.py

@@ -2,32 +2,26 @@ from __future__ import annotations
 
 from typing import Protocol
 
-from goose._internal.flow import FlowRun
-from goose._internal.state import FlowArguments, SerializedFlowRun
+from goose._internal.state import SerializedFlowRun
 
 
-class IFlowRunStore[FlowArgumentsT: FlowArguments](Protocol):
+class IFlowRunStore(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 get(self, *, run_id: str) -> SerializedFlowRun | None: ...
+    async def save(self, *, run_id: str, run: SerializedFlowRun) -> 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:
+class InMemoryFlowRunStore(IFlowRunStore):
+    def __init__(self, *, flow_name: str) -> 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 get(self, *, run_id: str) -> SerializedFlowRun | None:
+        return self._runs.get(run_id)
 
-    async def save(self, *, run: FlowRun[FlowArgumentsT]) -> None:
-        self._runs[run.id] = run.dump()
+    async def save(self, *, run_id: str, run: SerializedFlowRun) -> None:
+        self._runs[run_id] = run
 
     async def delete(self, *, run_id: str) -> None:
         self._runs.pop(run_id, None)

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

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

@@ -6,13 +6,13 @@ 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/flow.py,sha256=05U2f5i8ofQWDjghhomwuuEPMk-ftzXn7BVl_s7pIf8,4203
 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/state.py,sha256=LmMxhr3OrilYRHRxw1xYnl_dUkndaayLej83ccjfOGs,6583
+goose/_internal/store.py,sha256=GMW0wBpxESmRBLfL_lFKEi9x2P6Wd6-gZ7AWjWBTUmA,904
 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.9.4.dist-info/METADATA,sha256=awELFF40wuVHwy9Xl76Aaku36uLX2sxxFUOMbxwiN7c,441
+goose_py-0.9.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+goose_py-0.9.4.dist-info/RECORD,,