goose-py 0.3.12__py3-none-any.whl → 0.4.1__py3-none-any.whl

goose/agent.py

@@ -7,6 +7,7 @@ from typing import Any, ClassVar, Literal, NotRequired, Protocol, TypedDict
 
 from litellm import acompletion
 from pydantic import BaseModel, computed_field
+from goose.result import Result, TextResult
 
 
 class GeminiModel(StrEnum):
@@ -115,7 +116,7 @@ class AgentResponseDump(TypedDict):
     duration_ms: int
 
 
-class AgentResponse[R: BaseModel](BaseModel):
+class AgentResponse[R: BaseModel | str](BaseModel):
     INPUT_CENTS_PER_MILLION_TOKENS: ClassVar[dict[GeminiModel, float]] = {
         GeminiModel.FLASH_8B: 30,
         GeminiModel.FLASH: 15,
@@ -186,6 +187,12 @@ class AgentResponse[R: BaseModel](BaseModel):
             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,
             "flow_name": self.flow_name,
@@ -193,7 +200,7 @@ class AgentResponse[R: BaseModel](BaseModel):
             "model": self.model.value,
             "system_message": minimized_system_message,
             "input_messages": minimized_input_messages,
-            "output_message": self.response.model_dump_json(),
+            "output_message": output_message,
             "input_tokens": self.input_tokens,
             "output_tokens": self.output_tokens,
             "input_cost": self.input_cost,
@@ -221,13 +228,13 @@ class Agent:
         self.run_id = run_id
         self.logger = logger
 
-    async def __call__[R: BaseModel](
+    async def __call__[R: Result](
         self,
         *,
         messages: list[UserMessage | AssistantMessage],
         model: GeminiModel,
-        response_model: type[R],
         task_name: str,
+        response_model: type[R] = TextResult,
         system: SystemMessage | None = None,
     ) -> R:
         start_time = datetime.now()
@@ -235,22 +242,25 @@ class Agent:
         if system is not None:
             rendered_messages.insert(0, system.render())
 
-        response = await acompletion(
-            model=model.value,
-            messages=rendered_messages,
-            response_format={
-                "type": "json_object",
-                "response_schema": response_model.model_json_schema(),
-                "enforce_validation": True,
-            },
-        )
-
-        if len(response.choices) == 0:
-            raise RuntimeError("No content returned from LLM call.")
+        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}
+            )
+        else:
+            response = await acompletion(
+                model=model.value,
+                messages=rendered_messages,
+                response_format={
+                    "type": "json_object",
+                    "response_schema": response_model.model_json_schema(),
+                    "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(
             response=parsed_response,
@@ -271,4 +281,4 @@ class Agent:
         else:
             logging.info(agent_response.model_dump())
 
-        return agent_response.response
+        return parsed_response

goose/flow.py

@@ -13,7 +13,7 @@ from typing import (
     overload,
 )
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel
 
 from goose.agent import (
     Agent,
@@ -25,14 +25,11 @@ from goose.agent import (
 )
 from goose.errors import Honk
 from goose.store import IFlowRunStore, InMemoryFlowRunStore
+from goose.result import Result
 
 SerializedFlowRun = NewType("SerializedFlowRun", str)
 
 
-class Result(BaseModel):
-    model_config = ConfigDict(frozen=True)
-
-
 class Conversation[R: Result](BaseModel):
     user_messages: list[UserMessage]
     result_messages: list[R]
@@ -115,6 +112,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 +129,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()
-
-    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
+    @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")
 
-        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]] = []
@@ -166,6 +160,22 @@ class FlowRun:
                 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,
         *,
@@ -192,25 +202,35 @@ class FlowRun:
             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
 
 
@@ -264,8 +284,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__(
@@ -323,7 +356,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
 
@@ -331,7 +364,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/result.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class Result(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+
+class TextResult(Result):
+    text: str

{goose_py-0.3.12.dist-info → goose_py-0.4.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: goose-py
-Version: 0.3.12
+Version: 0.4.1
 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.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+goose/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+goose/agent.py,sha256=S9nDhK7v4Rsk-ANuq5tIZfZ7TOmCdMEVid2VuK1EWCw,8106
+goose/errors.py,sha256=-0OyZQJWYTRw5YgnCB2_uorVaUsL6Z0QYQO2FqzCiyg,32
+goose/flow.py,sha256=ZoMk3TlTLG1KM9ZxxR6lAuzDxzPZgUeWAi5mm4Ch7L4,13428
+goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+goose/result.py,sha256=-eZJn-2sPo7rHZ38Sz6IAHXqiJ-Ss39esEoFGimJEBI,155
+goose/store.py,sha256=4p2BBVAEUS1_Z0iBk5Qk_fPxRQeph64DRzXOFmjIT38,844
+goose_py-0.4.1.dist-info/METADATA,sha256=AB7TA_rXnj1M_gQ8Y_PnX7mhbdvJl7nlAAWf_S651fk,1106
+goose_py-0.4.1.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+goose_py-0.4.1.dist-info/RECORD,,

goose_py-0.3.12.dist-info/RECORD

@@ -1,9 +0,0 @@
-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=eZMaMQNhXqL5EMJLjDIRyqFKv9K_WV6sJsc50Oms1Vg,12152
-goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose/store.py,sha256=4p2BBVAEUS1_Z0iBk5Qk_fPxRQeph64DRzXOFmjIT38,844
-goose_py-0.3.12.dist-info/METADATA,sha256=ATNQi96Z1ZwwicyVMPOcM371cuU5mkEVXn4R-Jm_k-s,1107
-goose_py-0.3.12.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-goose_py-0.3.12.dist-info/RECORD,,