goose-py 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

goose/agent.py

@@ -1,8 +1,9 @@
 import base64
+import json
 import logging
 from datetime import datetime
 from enum import StrEnum
-from typing import Any, Awaitable, Callable, ClassVar, Literal, NotRequired, TypedDict
+from typing import Any, ClassVar, Literal, NotRequired, Protocol, TypedDict
 
 from litellm import acompletion
 from pydantic import BaseModel, computed_field
@@ -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,44 @@ 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: ...
+
 
 class Agent:
     def __init__(
@@ -150,7 +207,7 @@ class Agent:
         *,
         flow_name: str,
         run_id: str,
-        logger: Callable[[AgentResponse[Any]], Awaitable[None]] | None = None,
+        logger: IAgentLogger | None = None,
     ) -> None:
         self.flow_name = flow_name
         self.run_id = run_id
@@ -202,7 +259,7 @@ class Agent:
         )
 
         if self.logger is not None:
-            await self.logger(agent_response)
+            await self.logger(response=agent_response)
         else:
             logging.info(agent_response.model_dump())
 

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,
@@ -14,8 +14,16 @@ from typing import (
 
 from pydantic import BaseModel, ConfigDict, field_validator
 
-from goose.agent import Agent, AssistantMessage, LLMMessage, SystemMessage, UserMessage
+from goose.agent import (
+    Agent,
+    AssistantMessage,
+    IAgentLogger,
+    LLMMessage,
+    SystemMessage,
+    UserMessage,
+)
 from goose.errors import Honk
+from goose.store import IFlowRunStore, InMemoryFlowRunStore
 
 SerializedFlowRun = NewType("SerializedFlowRun", str)
 
@@ -174,11 +182,19 @@ class FlowRun:
                 last_input_hash=0,
             )
 
-    def start(self, *, flow_name: str, run_id: str) -> None:
+    def start(
+        self,
+        *,
+        flow_name: str,
+        run_id: str,
+        agent_logger: IAgentLogger | None = None,
+    ) -> None:
         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)
+        self._agent = Agent(
+            flow_name=self.flow_name, run_id=self.id, logger=agent_logger
+        )
 
     def end(self) -> None:
         self._last_requested_indices = {}
@@ -216,10 +232,18 @@ _current_flow_run: ContextVar[FlowRun | None] = ContextVar(
 
 class Flow[**P]:
     def __init__(
-        self, fn: Callable[P, Awaitable[None]], /, *, name: str | None = None
+        self,
+        fn: Callable[P, Awaitable[None]],
+        /,
+        *,
+        store: IFlowRunStore | None = None,
+        name: str | 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:
@@ -232,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)
+        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)
@@ -365,16 +389,20 @@ def task[**P, R: Result](
 def flow[**P](fn: Callable[P, Awaitable[None]], /) -> Flow[P]: ...
 @overload
 def flow[**P](
-    *, name: str | None = None
+    *, name: str | 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
+    fn: Callable[P, Awaitable[None]] | None = None,
+    /,
+    *,
+    name: str | 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)
+            return Flow(fn, name=name, agent_logger=agent_logger)
 
         return decorator
 
-    return Flow(fn, name=name)
+    return Flow(fn, name=name, 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.1.dist-info → goose_py-0.3.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: goose-py
-Version: 0.3.1
+Version: 0.3.3
 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.3.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=I7iWizu6MPHOJYJL-ZQyjbgvjxECBJekibkfO0bGKX0,12490
+goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+goose/store.py,sha256=6IIGkIYt6CWiNENi2gaS9C0Z-62ia-Qdz29HhCg8njw,698
+goose_py-0.3.3.dist-info/METADATA,sha256=KAvJyHSN3cfcjFUVQCRCGSpwGxKqk76JMkNWZRPCH_g,1106
+goose_py-0.3.3.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+goose_py-0.3.3.dist-info/RECORD,,

goose_py-0.3.1.dist-info/RECORD

@@ -1,8 +0,0 @@
-goose/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose/agent.py,sha256=cJO6Jlh8En_FGLzjelm0w1pM4rsVqfIOpfKUdpkord4,5650
-goose/errors.py,sha256=-0OyZQJWYTRw5YgnCB2_uorVaUsL6Z0QYQO2FqzCiyg,32
-goose/flow.py,sha256=_ws3Jrp96E5UG7D0dgT0_c1RDs07H-hgJ0mSAmJf_yE,11735
-goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose_py-0.3.1.dist-info/METADATA,sha256=wOP3QDcyVUJg-vjBw2KBb1xWdw_NA-50wgJtwg24WQc,1106
-goose_py-0.3.1.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-goose_py-0.3.1.dist-info/RECORD,,