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

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
@@ -112,7 +112,7 @@ class AgentResponse[R: BaseModel](BaseModel):
     }
 
     response: R
-    run_name: str
+    run_id: str
     flow_name: str
     task_name: str
     model: GeminiModel
@@ -144,16 +144,20 @@ 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_name: str,
-        logger: Callable[[AgentResponse[Any]], Awaitable[None]] | None = None,
+        run_id: str,
+        logger: IAgentLogger | None = None,
     ) -> None:
         self.flow_name = flow_name
-        self.run_name = run_name
+        self.run_id = run_id
         self.logger = logger
 
     async def __call__[R: BaseModel](
@@ -189,7 +193,7 @@ class Agent:
         end_time = datetime.now()
         agent_response = AgentResponse(
             response=parsed_response,
-            run_name=self.run_name,
+            run_id=self.run_id,
             flow_name=self.flow_name,
             task_name=task_name,
             model=model,
@@ -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/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)
@@ -122,12 +129,17 @@ class FlowRun:
     def __init__(self) -> None:
         self._node_states: dict[tuple[str, int], str] = {}
         self._last_requested_indices: dict[str, int] = {}
-        self._name = ""
+        self._flow_name = ""
+        self._id = ""
         self._agent: Agent | None = None
 
     @property
-    def name(self) -> str:
-        return self._name
+    def flow_name(self) -> str:
+        return self._flow_name
+
+    @property
+    def id(self) -> str:
+        return self._id
 
     @property
     def agent(self) -> Agent:
@@ -169,14 +181,24 @@ class FlowRun:
                 last_input_hash=0,
             )
 
-    def start(self, *, name: str) -> None:
+    def start(
+        self,
+        *,
+        flow_name: str,
+        run_id: str,
+        agent_logger: IAgentLogger | None = None,
+    ) -> None:
         self._last_requested_indices = {}
-        self._name = name
-        self._agent = Agent(flow_name=self.name, run_name=name)
+        self._flow_name = flow_name
+        self._id = run_id
+        self._agent = Agent(
+            flow_name=self.flow_name, run_id=self.id, logger=agent_logger
+        )
 
     def end(self) -> None:
         self._last_requested_indices = {}
-        self._name = ""
+        self._flow_name = ""
+        self._id = ""
         self._agent = None
 
     def dump(self) -> SerializedFlowRun:
@@ -209,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:
@@ -226,14 +254,18 @@ class Flow[**P]:
         return run
 
     @contextmanager
-    def start_run(self, *, name: str, run: FlowRun | None = None) -> Iterator[FlowRun]:
-        if run is None:
+    def start_run(
+        self, *, run_id: str, preload: FlowRun | None = None
+    ) -> Iterator[FlowRun]:
+        if preload is None:
             run = FlowRun()
+        else:
+            run = preload
 
         old_run = _current_flow_run.get()
         _current_flow_run.set(run)
 
-        run.start(name=name)
+        run.start(flow_name=self.name, run_id=run_id, agent_logger=self._agent_logger)
         yield run
         run.end()
 
@@ -354,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.0.dist-info → goose_py-0.3.2.dist-info}/METADATA

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

@@ -0,0 +1,8 @@
+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,,

goose_py-0.3.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-goose/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose/agent.py,sha256=j8r9p7PqIdXCX_Vl8pUTrF2SshJXUDitgJlCV_C6hdw,5662
-goose/errors.py,sha256=-0OyZQJWYTRw5YgnCB2_uorVaUsL6Z0QYQO2FqzCiyg,32
-goose/flow.py,sha256=5a9lIHTyW5hLeMdt1AfIxe5z-VlTti_3OJ1o4f_nuoo,11458
-goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose_py-0.3.0.dist-info/METADATA,sha256=dMNVLzZZ7KyEUacFGWqLeNJuv0ni5-GS_FP9ifaF4bM,1106
-goose_py-0.3.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-goose_py-0.3.0.dist-info/RECORD,,