goose-py 0.3.1__tar.gz → 0.3.2__tar.gz

{goose_py-0.3.1 → goose_py-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: goose-py
-Version: 0.3.1
+Version: 0.3.2
 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.1 → goose_py-0.3.2}/goose/agent.py

@@ -2,7 +2,7 @@ import base64
 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
@@ -144,13 +144,17 @@ class AgentResponse[R: BaseModel](BaseModel):
         return self.input_cost + self.output_cost
 
 
+class IAgentLogger(Protocol):
+    async def __call__(self, *, response: AgentResponse[Any]) -> None: ...
+
+
 class Agent:
     def __init__(
         self,
         *,
         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 +206,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_py-0.3.1 → goose_py-0.3.2}/goose/flow.py

@@ -14,7 +14,14 @@ 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
 
 SerializedFlowRun = NewType("SerializedFlowRun", str)
@@ -174,11 +181,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 +231,16 @@ _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]],
+        /,
+        *,
+        name: str | None = None,
+        agent_logger: IAgentLogger | None = None,
     ) -> None:
         self._fn = fn
         self._name = name
+        self._agent_logger = agent_logger
 
     @property
     def name(self) -> str:
@@ -244,7 +265,7 @@ class Flow[**P]:
         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
         run.end()
 
@@ -365,16 +386,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_py-0.3.1 → goose_py-0.3.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "goose-py"
-version = "0.3.1"
+version = "0.3.2"
 description = "A tool for AI workflows based on human-computer collaboration and structured output."
 authors = [
     "Nash Taylor <nash@chelle.ai>",