goose-py 0.7.0__py3-none-any.whl → 0.7.2__py3-none-any.whl

goose/_internal/agent.py

@@ -65,20 +65,12 @@ class AgentResponse[R: BaseModel | str](BaseModel):
     @computed_field
     @property
     def input_cost(self) -> float:
-        return (
-            self.INPUT_CENTS_PER_MILLION_TOKENS[self.model]
-            * self.input_tokens
-            / 1_000_000
-        )
+        return self.INPUT_CENTS_PER_MILLION_TOKENS[self.model] * self.input_tokens / 1_000_000
 
     @computed_field
     @property
     def output_cost(self) -> float:
-        return (
-            self.OUTPUT_CENTS_PER_MILLION_TOKENS[self.model]
-            * self.output_tokens
-            / 1_000_000
-        )
+        return self.OUTPUT_CENTS_PER_MILLION_TOKENS[self.model] * self.output_tokens / 1_000_000
 
     @computed_field
     @property
@@ -100,15 +92,9 @@ class AgentResponse[R: BaseModel | str](BaseModel):
             for part in message["content"]:
                 if part["type"] == "image_url":
                     part["image_url"] = "__MEDIA__"
-        minimized_input_messages = [
-            json.dumps(message) for message in minimized_input_messages
-        ]
-
-        output_message = (
-            self.response.model_dump_json()
-            if isinstance(self.response, BaseModel)
-            else self.response
-        )
+        minimized_input_messages = [json.dumps(message) for message in minimized_input_messages]
+
+        output_message = self.response.model_dump_json() if isinstance(self.response, BaseModel) else self.response
 
         return {
             "run_id": self.run_id,
@@ -161,9 +147,7 @@ class Agent:
 
         if response_model is TextResult:
             response = await acompletion(model=model.value, messages=rendered_messages)
-            parsed_response = response_model.model_validate(
-                {"text": response.choices[0].message.content}
-            )
+            parsed_response = response_model.model_validate({"text": response.choices[0].message.content})
         else:
             response = await acompletion(
                 model=model.value,
@@ -174,9 +158,7 @@ class Agent:
                     "enforce_validation": True,
                 },
             )
-            parsed_response = response_model.model_validate_json(
-                response.choices[0].message.content
-            )
+            parsed_response = response_model.model_validate_json(response.choices[0].message.content)
 
         end_time = datetime.now()
         agent_response = AgentResponse(

goose/_internal/conversation.py

@@ -1,12 +1,9 @@
+from typing import Self
+
 from pydantic import BaseModel
 
 from goose._internal.result import Result
-from goose._internal.types.agent import (
-    AssistantMessage,
-    LLMMessage,
-    SystemMessage,
-    UserMessage,
-)
+from goose._internal.types.agent import AssistantMessage, LLMMessage, SystemMessage, UserMessage
 
 
 class Conversation[R: Result](BaseModel):
@@ -24,18 +21,15 @@ class Conversation[R: Result](BaseModel):
             messages.append(self.context.render())
 
         for message_index in range(len(self.user_messages)):
-            messages.append(
-                AssistantMessage(
-                    text=self.result_messages[message_index].model_dump_json()
-                ).render()
-            )
+            messages.append(AssistantMessage(text=self.result_messages[message_index].model_dump_json()).render())
             messages.append(self.user_messages[message_index].render())
 
         if len(self.result_messages) > len(self.user_messages):
-            messages.append(
-                AssistantMessage(
-                    text=self.result_messages[-1].model_dump_json()
-                ).render()
-            )
+            messages.append(AssistantMessage(text=self.result_messages[-1].model_dump_json()).render())
 
         return messages
+
+    def undo(self) -> Self:
+        self.user_messages.pop()
+        self.result_messages.pop()
+        return self

goose/_internal/flow.py

@@ -1,6 +1,7 @@
+from collections.abc import AsyncIterator, Awaitable, Callable
 from contextlib import asynccontextmanager
 from types import CodeType
-from typing import AsyncIterator, Awaitable, Callable, Protocol, overload
+from typing import Protocol, overload
 
 from goose._internal.agent import Agent, IAgentLogger
 from goose._internal.conversation import Conversation
@@ -13,9 +14,7 @@ from goose.errors import Honk
 class IAdapter[ResultT: Result](Protocol):
     __code__: CodeType
 
-    async def __call__(
-        self, *, conversation: Conversation[ResultT], agent: Agent
-    ) -> ResultT: ...
+    async def __call__(self, *, conversation: Conversation[ResultT], agent: Agent) -> ResultT: ...
 
 
 class Flow[**P]:

goose/_internal/state.py

@@ -61,6 +61,10 @@ class NodeState[ResultT: Result](BaseModel):
         self.conversation.user_messages.append(message)
         return self
 
+    def undo(self) -> Self:
+        self.conversation.undo()
+        return self
+
 
 class FlowRun:
     def __init__(self) -> None:
@@ -97,23 +101,17 @@ class FlowRun:
         matching_nodes: list[NodeState[R]] = []
         for key, node_state in self._node_states.items():
             if key[0] == task.name:
-                matching_nodes.append(
-                    NodeState[task.result_type].model_validate_json(node_state)
-                )
+                matching_nodes.append(NodeState[task.result_type].model_validate_json(node_state))
         return sorted(matching_nodes, key=lambda node: node.index)
 
     def get[R: Result](self, *, task: "Task[Any, R]", index: int = 0) -> NodeState[R]:
-        if (
-            existing_node_state := self._node_states.get((task.name, index))
-        ) is not None:
+        if (existing_node_state := self._node_states.get((task.name, index))) is not None:
             return NodeState[task.result_type].model_validate_json(existing_node_state)
         else:
             return NodeState[task.result_type](
                 task_name=task.name,
                 index=index,
-                conversation=Conversation[task.result_type](
-                    user_messages=[], result_messages=[]
-                ),
+                conversation=Conversation[task.result_type](user_messages=[], result_messages=[]),
                 last_hash=0,
             )
 
@@ -143,9 +141,7 @@ class FlowRun:
         self._last_requested_indices = {}
         self._flow_name = flow_name
         self._id = run_id
-        self._agent = Agent(
-            flow_name=self.flow_name, run_id=self.id, logger=agent_logger
-        )
+        self._agent = Agent(flow_name=self.flow_name, run_id=self.id, logger=agent_logger)
 
     def end(self) -> None:
         self._last_requested_indices = {}
@@ -177,9 +173,7 @@ class FlowRun:
         return flow_run
 
 
-_current_flow_run: ContextVar[FlowRun | None] = ContextVar(
-    "current_flow_run", default=None
-)
+_current_flow_run: ContextVar[FlowRun | None] = ContextVar("current_flow_run", default=None)
 
 
 def get_current_flow_run() -> FlowRun | None:

goose/_internal/task.py

@@ -1,4 +1,5 @@
-from typing import Awaitable, Callable, overload
+from collections.abc import Awaitable, Callable
+from typing import overload
 
 from goose._internal.agent import Agent, GeminiModel, SystemMessage, UserMessage
 from goose._internal.conversation import Conversation
@@ -33,9 +34,7 @@ class Task[**P, R: Result]:
     def name(self) -> str:
         return self._generator.__name__
 
-    async def generate(
-        self, state: NodeState[R], *args: P.args, **kwargs: P.kwargs
-    ) -> R:
+    async def generate(self, state: NodeState[R], *args: P.args, **kwargs: P.kwargs) -> R:
         state_hash = self.__hash_task_call(*args, **kwargs)
         if state_hash != state.last_hash:
             result = await self._generator(*args, **kwargs)
@@ -58,14 +57,18 @@ class Task[**P, R: Result]:
             node_state.set_context(context=context)
         node_state.add_user_message(message=user_message)
 
-        result = await self.__adapt(
-            conversation=node_state.conversation, agent=flow_run.agent
-        )
+        result = await self.__adapt(conversation=node_state.conversation, agent=flow_run.agent)
         node_state.add_result(result=result)
         flow_run.add_node_state(node_state)
 
         return result
 
+    def undo(self, *, index: int = 0) -> None:
+        flow_run = self.__get_current_flow_run()
+        node_state = flow_run.get(task=self, index=index)
+        node_state.undo()
+        flow_run.add_node_state(node_state)
+
     async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
         flow_run = self.__get_current_flow_run()
         node_state = flow_run.get_next(task=self)
@@ -97,11 +100,7 @@ class Task[**P, R: Result]:
 
     def __hash_task_call(self, *args: P.args, **kwargs: P.kwargs) -> int:
         try:
-            to_hash = str(
-                tuple(args)
-                + tuple(kwargs.values())
-                + (self._generator.__code__, self._adapter_model)
-            )
+            to_hash = str(tuple(args) + tuple(kwargs.values()) + (self._generator.__code__, self._adapter_model))
             return hash(to_hash)
         except TypeError:
             raise Honk(f"Unhashable argument to task {self.name}: {args} {kwargs}")

{goose_py-0.7.0.dist-info → goose_py-0.7.2.dist-info}/METADATA

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

@@ -4,15 +4,15 @@ 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=lXx7tXS518Od7ASIAcFnaGxb_GyJjvm2DI-ufvPJqa8,6101
-goose/_internal/conversation.py,sha256=lOJ3v4uXjM8l8_kRShFTPJiRo4bJr7JvCEl8gwOdres,1217
-goose/_internal/flow.py,sha256=334uZEkgthWqe2kNloArnIPWUuY9TXSdfItrr82Yn7E,3253
+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=3psalnzNN7lZkc5jKGqprCSjt31z6ZiY3wzcefpa-Go,5829
+goose/_internal/state.py,sha256=NTF9L2hrO7xXKPNNJjmOhHiXSTW5mnDULjvm81aqGLQ,5785
 goose/_internal/store.py,sha256=vIxPIpechF_lEQlQ8JT1NDySDfHe3-eMHEWeTqVbscg,946
-goose/_internal/task.py,sha256=Sru14gAoebW5LXWWCJxs_HWiivhO-zq8cXr6bOosAcA,4756
+goose/_internal/task.py,sha256=k_DuFTnaugbkFXz_xvZOrmwFV-MlKf0RZCXfR7M_6ko,4910
 goose/_internal/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 goose/_internal/types/agent.py,sha256=rNVt2gEr_m4_8tGFgcdichpPp8xhOS5GY0kN2C4tiE8,2153
-goose_py-0.7.0.dist-info/METADATA,sha256=Petgz53R3syhL8FAuSe32QUEMrtGFg_hNvCbC2ZoI7w,441
-goose_py-0.7.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-goose_py-0.7.0.dist-info/RECORD,,
+goose_py-0.7.2.dist-info/METADATA,sha256=FVAKqtk6EqumBkzOEtbFn9GA1mLiAinNF-nJ7-K6dC0,441
+goose_py-0.7.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+goose_py-0.7.2.dist-info/RECORD,,