goose-py 0.3.11__py3-none-any.whl → 0.4.0__py3-none-any.whl

goose/agent.py

@@ -10,10 +10,9 @@ from pydantic import BaseModel, computed_field
 
 
 class GeminiModel(StrEnum):
-    EXP = "gemini/gemini-exp-1121"
-    PRO = "gemini/gemini-1.5-pro"
-    FLASH = "gemini/gemini-1.5-flash"
-    FLASH_8B = "gemini/gemini-1.5-flash-8b"
+    PRO = "vertex_ai/gemini-1.5-pro"
+    FLASH = "vertex_ai/gemini-1.5-flash"
+    FLASH_8B = "vertex_ai/gemini-1.5-flash-8b"
 
 
 class UserMediaContentType(StrEnum):
@@ -121,13 +120,11 @@ class AgentResponse[R: BaseModel](BaseModel):
         GeminiModel.FLASH_8B: 30,
         GeminiModel.FLASH: 15,
         GeminiModel.PRO: 500,
-        GeminiModel.EXP: 0,
     }
     OUTPUT_CENTS_PER_MILLION_TOKENS: ClassVar[dict[GeminiModel, float]] = {
         GeminiModel.FLASH_8B: 30,
         GeminiModel.FLASH: 15,
         GeminiModel.PRO: 500,
-        GeminiModel.EXP: 0,
     }
 
     response: R

goose/flow.py

@@ -115,6 +115,8 @@ class FlowRun:
         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
 
     @property
     def flow_name(self) -> str:
@@ -130,17 +132,12 @@ class FlowRun:
             raise Honk("Agent is only accessible once a run is started")
         return self._agent
 
-    def add(self, node_state: NodeState[Any], /) -> None:
-        key = (node_state.task_name, node_state.index)
-        self._node_states[key] = node_state.model_dump_json()
+    @property
+    def flow_inputs(self) -> tuple[tuple[Any, ...], dict[str, Any]]:
+        if self._flow_args is None or self._flow_kwargs is None:
+            raise Honk("This Flow run has not been executed before")
 
-    def get_next[R: Result](self, *, task: "Task[Any, R]") -> NodeState[R]:
-        if task.name not in self._last_requested_indices:
-            self._last_requested_indices[task.name] = 0
-        else:
-            self._last_requested_indices[task.name] += 1
-
-        return self.get(task=task, index=self._last_requested_indices[task.name])
+        return self._flow_args, self._flow_kwargs
 
     def get_all[R: Result](self, *, task: "Task[Any, R]") -> list[NodeState[R]]:
         matching_nodes: list[NodeState[R]] = []
@@ -159,13 +156,29 @@ class FlowRun:
         else:
             return NodeState[task.result_type](
                 task_name=task.name,
-                index=index or 0,
+                index=index,
                 conversation=Conversation[task.result_type](
                     user_messages=[], result_messages=[]
                 ),
                 last_hash=0,
             )
 
+    def set_flow_inputs(self, *args: Any, **kwargs: Any) -> None:
+        self._flow_args = args
+        self._flow_kwargs = kwargs
+
+    def add_node_state(self, node_state: NodeState[Any], /) -> None:
+        key = (node_state.task_name, node_state.index)
+        self._node_states[key] = node_state.model_dump_json()
+
+    def get_next[R: Result](self, *, task: "Task[Any, R]") -> NodeState[R]:
+        if task.name not in self._last_requested_indices:
+            self._last_requested_indices[task.name] = 0
+        else:
+            self._last_requested_indices[task.name] += 1
+
+        return self.get(task=task, index=self._last_requested_indices[task.name])
+
     def start(
         self,
         *,
@@ -186,26 +199,41 @@ class FlowRun:
         self._id = ""
         self._agent = None
 
+    def clear_node(self, *, task: "Task[Any, Result]", index: int) -> None:
+        key = (task.name, index)
+        if key in self._node_states:
+            del self._node_states[key]
+
     def dump(self) -> SerializedFlowRun:
+        flow_args, flow_kwargs = self.flow_inputs
+
         return SerializedFlowRun(
             json.dumps(
                 {
-                    ":".join([task_name, str(index)]): value
-                    for (task_name, index), value in self._node_states.items()
+                    "node_states": {
+                        ":".join([task_name, str(index)]): value
+                        for (task_name, index), value in self._node_states.items()
+                    },
+                    "flow_args": list(flow_args),
+                    "flow_kwargs": flow_kwargs,
                 }
             )
         )
 
     @classmethod
-    def load(cls, run: SerializedFlowRun, /) -> Self:
+    def load(cls, serialized_flow_run: SerializedFlowRun, /) -> Self:
         flow_run = cls()
-        raw_node_states = json.loads(run)
+        run = json.loads(serialized_flow_run)
+
         new_node_states: dict[tuple[str, int], str] = {}
-        for key, node_state in raw_node_states.items():
+        for key, node_state in run["node_states"].items():
             task_name, index = tuple(key.split(":"))
             new_node_states[(task_name, int(index))] = node_state
-
         flow_run._node_states = new_node_states
+
+        flow_run._flow_args = tuple(run["flow_args"])
+        flow_run._flow_kwargs = run["flow_kwargs"]
+
         return flow_run
 
 
@@ -259,8 +287,21 @@ class Flow[**P]:
         _current_flow_run.set(old_run)
 
     async def generate(self, *args: P.args, **kwargs: P.kwargs) -> None:
+        flow_run = _current_flow_run.get()
+        if flow_run is None:
+            raise Honk("No current flow run")
+
+        flow_run.set_flow_inputs(*args, **kwargs)
         await self._fn(*args, **kwargs)
 
+    async def regenerate(self) -> None:
+        flow_run = _current_flow_run.get()
+        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)
+
 
 class Task[**P, R: Result]:
     def __init__(
@@ -318,7 +359,7 @@ class Task[**P, R: Result]:
 
         result = await self._adapter(conversation=node_state.conversation)
         node_state.add_result(result=result)
-        flow_run.add(node_state)
+        flow_run.add_node_state(node_state)
 
         return result
 
@@ -326,7 +367,7 @@ class Task[**P, R: Result]:
         flow_run = self.__get_current_flow_run()
         node_state = flow_run.get_next(task=self)
         result = await self.generate(node_state, *args, **kwargs)
-        flow_run.add(node_state)
+        flow_run.add_node_state(node_state)
         return result
 
     def __hash_task_call(self, *args: P.args, **kwargs: P.kwargs) -> int:

goose/store.py

@@ -10,6 +10,7 @@ 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):
@@ -22,3 +23,6 @@ class InMemoryFlowRunStore(IFlowRunStore):
 
     async def save(self, *, run: FlowRun) -> None:
         self._runs[run.id] = run
+
+    async def delete(self, *, run_id: str) -> None:
+        self._runs.pop(run_id, None)

{goose_py-0.3.11.dist-info → goose_py-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: goose-py
-Version: 0.3.11
+Version: 0.4.0
 Summary: A tool for AI workflows based on human-computer collaboration and structured output.
 Home-page: https://github.com/chelle-ai/goose
 Keywords: ai,yaml,configuration,llm

goose_py-0.4.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+goose/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+goose/agent.py,sha256=pZkKKW2weXEX43p9-pVIpEtwei5nylyQ_YrNZWYFQM0,7687
+goose/errors.py,sha256=-0OyZQJWYTRw5YgnCB2_uorVaUsL6Z0QYQO2FqzCiyg,32
+goose/flow.py,sha256=tzubasKLCOgB_X5-MgC-44OhbwqFZq_7gja3sSElCMI,13478
+goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+goose/store.py,sha256=4p2BBVAEUS1_Z0iBk5Qk_fPxRQeph64DRzXOFmjIT38,844
+goose_py-0.4.0.dist-info/METADATA,sha256=CJhAf8NiFEOq7pL4Ng4zQbuLMe9Gax-cO1UXJHnYpLY,1106
+goose_py-0.4.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+goose_py-0.4.0.dist-info/RECORD,,

goose_py-0.3.11.dist-info/RECORD

@@ -1,9 +0,0 @@
-goose/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose/agent.py,sha256=AH_yJF3i79cFz4nHamgkXXZy2KgNbmiVNsolZaB5vO8,7769
-goose/errors.py,sha256=-0OyZQJWYTRw5YgnCB2_uorVaUsL6Z0QYQO2FqzCiyg,32
-goose/flow.py,sha256=Oyq3hVe3Yhs5SqC1sNFl0pn1fZ6qgGkoClwzuaAlAPk,11971
-goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose/store.py,sha256=6IIGkIYt6CWiNENi2gaS9C0Z-62ia-Qdz29HhCg8njw,698
-goose_py-0.3.11.dist-info/METADATA,sha256=7qwWawDrI0qDYgkA_LT_2-AxUfGVnMXCgbktQ_bq5T8,1107
-goose_py-0.3.11.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-goose_py-0.3.11.dist-info/RECORD,,