goose-py 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

goose/agent.py

@@ -1,4 +1,5 @@
 import base64
+import json
 import logging
 from datetime import datetime
 from enum import StrEnum
@@ -97,6 +98,24 @@ class SystemMessage(BaseModel):
         }
 
 
+class AgentResponseDump(TypedDict):
+    run_id: str
+    flow_name: str
+    task_name: str
+    model: str
+    system_message: str
+    input_messages: str
+    output_message: str
+    input_cost: float
+    output_cost: float
+    total_cost: float
+    input_tokens: int
+    output_tokens: int
+    start_time: datetime
+    end_time: datetime
+    duration_ms: int
+
+
 class AgentResponse[R: BaseModel](BaseModel):
     INPUT_CENTS_PER_MILLION_TOKENS: ClassVar[dict[GeminiModel, float]] = {
         GeminiModel.FLASH_8B: 30,
@@ -143,6 +162,40 @@ class AgentResponse[R: BaseModel](BaseModel):
     def total_cost(self) -> float:
         return self.input_cost + self.output_cost
 
+    def minimized_dump(self) -> AgentResponseDump:
+        if self.system is None:
+            minimized_system_message = ""
+        else:
+            minimized_system_message = self.system.model_dump()
+            for part in minimized_system_message["parts"]:
+                if part["type"] == "image_url":
+                    part["content"] = b"__MEDIA__"
+
+        minimized_input_messages = [
+            message.model_dump() for message in self.input_messages
+        ]
+        for message in minimized_input_messages:
+            if message["type"] == "image_url":
+                message["content"] = b"__MEDIA__"
+
+        return {
+            "run_id": self.run_id,
+            "flow_name": self.flow_name,
+            "task_name": self.task_name,
+            "model": self.model.value,
+            "system_message": json.dumps(minimized_system_message),
+            "input_messages": json.dumps(minimized_input_messages),
+            "output_message": self.response.model_dump_json(),
+            "input_tokens": self.input_tokens,
+            "output_tokens": self.output_tokens,
+            "input_cost": self.input_cost,
+            "output_cost": self.output_cost,
+            "total_cost": self.total_cost,
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "duration_ms": self.duration_ms,
+        }
+
 
 class IAgentLogger(Protocol):
     async def __call__(self, *, response: AgentResponse[Any]) -> None: ...

goose/flow.py

@@ -1,11 +1,11 @@
 import json
-from contextlib import contextmanager
+from contextlib import asynccontextmanager
 from contextvars import ContextVar
 from typing import (
     Any,
+    AsyncIterator,
     Awaitable,
     Callable,
-    Iterator,
     NewType,
     Protocol,
     Self,
@@ -23,6 +23,7 @@ from goose.agent import (
     UserMessage,
 )
 from goose.errors import Honk
+from goose.store import IFlowRunStore, InMemoryFlowRunStore
 
 SerializedFlowRun = NewType("SerializedFlowRun", str)
 
@@ -236,11 +237,13 @@ class Flow[**P]:
         /,
         *,
         name: str | 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)
 
     @property
     def name(self) -> str:
@@ -253,20 +256,20 @@ class Flow[**P]:
             raise Honk("No current flow run")
         return run
 
-    @contextmanager
-    def start_run(
-        self, *, run_id: str, preload: FlowRun | None = None
-    ) -> Iterator[FlowRun]:
-        if preload is None:
+    @asynccontextmanager
+    async def start_run(self, *, run_id: str) -> AsyncIterator[FlowRun]:
+        existing_run = await self._store.get(run_id=run_id)
+        if existing_run is None:
             run = FlowRun()
         else:
-            run = preload
+            run = existing_run
 
         old_run = _current_flow_run.get()
         _current_flow_run.set(run)
 
         run.start(flow_name=self.name, run_id=run_id, agent_logger=self._agent_logger)
         yield run
+        await self._store.save(run=run)
         run.end()
 
         _current_flow_run.set(old_run)
@@ -386,20 +389,24 @@ def task[**P, R: Result](
 def flow[**P](fn: Callable[P, Awaitable[None]], /) -> Flow[P]: ...
 @overload
 def flow[**P](
-    *, name: str | None = None, agent_logger: IAgentLogger | None = None
+    *,
+    name: str | None = None,
+    store: IFlowRunStore | None = None,
+    agent_logger: IAgentLogger | None = None,
 ) -> Callable[[Callable[P, Awaitable[None]]], Flow[P]]: ...
 def flow[**P](
     fn: Callable[P, Awaitable[None]] | None = None,
     /,
     *,
     name: str | None = None,
+    store: IFlowRunStore | None = None,
     agent_logger: IAgentLogger | None = None,
 ) -> Flow[P] | Callable[[Callable[P, Awaitable[None]]], Flow[P]]:
     if fn is None:
 
         def decorator(fn: Callable[P, Awaitable[None]]) -> Flow[P]:
-            return Flow(fn, name=name, agent_logger=agent_logger)
+            return Flow(fn, name=name, store=store, agent_logger=agent_logger)
 
         return decorator
 
-    return Flow(fn, name=name, agent_logger=agent_logger)
+    return Flow(fn, name=name, store=store, agent_logger=agent_logger)

goose/store.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    from goose.flow import FlowRun
+
+
+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: ...
+
+
+class InMemoryFlowRunStore(IFlowRunStore):
+    def __init__(self, *, flow_name: str) -> None:
+        self._flow_name = flow_name
+        self._runs: dict[str, FlowRun] = {}
+
+    async def get(self, *, run_id: str) -> FlowRun | None:
+        return self._runs.get(run_id)
+
+    async def save(self, *, run: FlowRun) -> None:
+        self._runs[run.id] = run

{goose_py-0.3.2.dist-info → goose_py-0.3.4.dist-info}/METADATA

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

@@ -0,0 +1,9 @@
+goose/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+goose/agent.py,sha256=uCqHl5BKUX5niPbuZBOqtwnmHf8IZeBcOMylVLUfQlY,7453
+goose/errors.py,sha256=-0OyZQJWYTRw5YgnCB2_uorVaUsL6Z0QYQO2FqzCiyg,32
+goose/flow.py,sha256=_kSNB9CnDoduOm0YLLq0q5kzkUrtNEtvAx7A9D0_YkU,12605
+goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+goose/store.py,sha256=6IIGkIYt6CWiNENi2gaS9C0Z-62ia-Qdz29HhCg8njw,698
+goose_py-0.3.4.dist-info/METADATA,sha256=XvAPlMqFJIZ4TLmSwPllzxcylb8D3w8PTi0qO0Y8v4c,1106
+goose_py-0.3.4.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+goose_py-0.3.4.dist-info/RECORD,,

goose_py-0.3.2.dist-info/RECORD

@@ -1,8 +0,0 @@
-goose/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose/agent.py,sha256=7QMlyCGezyRmNzG6FAjhHZcI00VFKJP-991ZpUEQRyg,5720
-goose/errors.py,sha256=-0OyZQJWYTRw5YgnCB2_uorVaUsL6Z0QYQO2FqzCiyg,32
-goose/flow.py,sha256=EHVnHLSMobzPto7Ydviv6sE2vKvmJfC0BLSpOa7GyAQ,12223
-goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose_py-0.3.2.dist-info/METADATA,sha256=JPBa5_zh5AAGZDmGxAvnI9fKQQu7pM88ppYDXfDs9No,1106
-goose_py-0.3.2.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-goose_py-0.3.2.dist-info/RECORD,,