goose-py 0.1.4__tar.gz → 0.2.1__tar.gz

{goose_py-0.1.4 → goose_py-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: goose-py
-Version: 0.1.4
+Version: 0.2.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
@@ -25,7 +25,7 @@ Project-URL: Documentation, https://github.com/chelle-ai/goose
 Project-URL: Repository, https://github.com/chelle-ai/goose
 Description-Content-Type: text/markdown
 
-# Vers
+# Goose
 
 Docs to come.
 

{goose_py-0.1.4 → goose_py-0.2.1}/README.md

@@ -1,3 +1,3 @@
-# Vers
+# Goose
 
 Docs to come.

goose_py-0.2.1/goose/__init__.py

@@ -0,0 +1,3 @@
+from goose.flow import ConversationState, FlowState, Result, flow, task
+
+__all__ = ["ConversationState", "FlowState", "Result", "flow", "task"]

goose_py-0.1.4/goose/types.py → goose_py-0.2.1/goose/agent.py

@@ -1,8 +1,10 @@
 import base64
+import logging
 from datetime import datetime
 from enum import StrEnum
-from typing import ClassVar, Literal, NotRequired, TypedDict
+from typing import Any, Awaitable, Callable, ClassVar, Literal, NotRequired, TypedDict
 
+from litellm import acompletion
 from pydantic import BaseModel, computed_field
 
 
@@ -131,3 +133,68 @@ class AgentResponse[R: BaseModel](BaseModel):
             self.OUTPUT_CENTS_PER_MILLION_TOKENS[self.model] * self.output_tokens
         )
         return input_cost + output_cost
+
+
+class Agent:
+    def __init__(
+        self,
+        *,
+        flow_name: str,
+        run_id: str,
+        logger: Callable[[AgentResponse[Any]], Awaitable[None]] | None = None,
+    ) -> None:
+        self.flow_name = flow_name
+        self.run_id = run_id
+        self.logger = logger
+
+    async def __call__[R: BaseModel](
+        self,
+        *,
+        messages: list[UserMessage | AssistantMessage],
+        model: GeminiModel,
+        response_model: type[R],
+        task_name: str,
+        system: SystemMessage | None = None,
+    ) -> R:
+        start_time = datetime.now()
+        rendered_messages = [message.render() for message in messages]
+        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.")
+
+        parsed_response = response_model.model_validate_json(
+            response.choices[0].message.content
+        )
+        end_time = datetime.now()
+        agent_response = AgentResponse(
+            response=parsed_response,
+            id=self.run_id,
+            flow_name=self.flow_name,
+            task_name=task_name,
+            model=model,
+            system=system,
+            input_messages=messages,
+            input_tokens=response.usage.prompt_tokens,
+            output_tokens=response.usage.completion_tokens,
+            start_time=start_time,
+            end_time=end_time,
+        )
+
+        if self.logger is not None:
+            await self.logger(agent_response)
+        else:
+            logging.info(agent_response.model_dump())
+
+        return agent_response.response

goose_py-0.2.1/goose/flow.py

@@ -0,0 +1,334 @@
+import json
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Iterator,
+    NewType,
+    Protocol,
+    Self,
+    overload,
+)
+
+from pydantic import BaseModel, ConfigDict, field_validator
+
+from goose.agent import UserMessage
+from goose.errors import Honk
+
+SerializedFlowState = NewType("SerializedFlowState", str)
+
+
+class Result(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+
+class GooseResponse[R: Result](BaseModel):
+    result: R
+
+
+class ConversationState[R: Result](BaseModel):
+    messages: list[UserMessage | GooseResponse[R]]
+
+    @field_validator("messages")
+    def alternates_starting_with_result(
+        cls, messages: list[UserMessage | GooseResponse[R]]
+    ) -> list[UserMessage | GooseResponse[R]]:
+        if len(messages) == 0:
+            return messages
+        elif isinstance(messages[0], UserMessage):
+            raise Honk(
+                "User cannot start a conversation on a Task, must begin with a Result"
+            )
+
+        last_message_type: type[UserMessage | GooseResponse[R]] = type(messages[0])
+        for message in messages[1:]:
+            if isinstance(message, last_message_type):
+                raise Honk(
+                    "Conversation must alternate between User and Result messages"
+                )
+            last_message_type = type(message)
+
+        return messages
+
+    @property
+    def awaiting_response(self) -> bool:
+        return len(self.messages) % 2 == 0
+
+
+class IAdapter[ResultT: Result](Protocol):
+    async def __call__(
+        self, *, conversation_state: ConversationState[ResultT]
+    ) -> ResultT: ...
+
+
+class NodeState[ResultT: Result](BaseModel):
+    task_name: str
+    index: int
+    conversation_state: ConversationState[ResultT]
+    last_input_hash: int
+    pinned: bool
+
+    @property
+    def result(self) -> ResultT:
+        last_message = self.conversation_state.messages[-1]
+        if isinstance(last_message, GooseResponse):
+            return last_message.result
+        else:
+            raise Honk("Node awaiting response, has no result")
+
+    def add_result(
+        self,
+        *,
+        result: ResultT,
+        new_input_hash: int | None = None,
+        overwrite: bool = False,
+    ) -> Self:
+        if overwrite:
+            if len(self.conversation_state.messages) == 0:
+                self.conversation_state.messages.append(GooseResponse(result=result))
+            else:
+                self.conversation_state.messages[-1] = GooseResponse(result=result)
+        else:
+            self.conversation_state.messages.append(GooseResponse(result=result))
+        if new_input_hash is not None:
+            self.last_input_hash = new_input_hash
+        return self
+
+    def add_user_message(self, *, message: UserMessage) -> Self:
+        self.conversation_state.messages.append(message)
+        return self
+
+    def pin(self) -> Self:
+        self.pinned = True
+        return self
+
+    def unpin(self) -> Self:
+        self.pinned = False
+        return self
+
+
+class FlowState:
+    def __init__(self) -> None:
+        self._node_states: dict[tuple[str, int], str] = {}
+        self._last_requested_indices: dict[str, int] = {}
+
+    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
+
+        return self.get(task=task, index=self._last_requested_indices[task.name])
+
+    def get_all[R: Result](self, *, task: "Task[Any, R]") -> list[NodeState[R]]:
+        matching_nodes: list[NodeState[R]] = []
+        for key, node_state in self._node_states.items():
+            if key[0] == task.name:
+                matching_nodes.append(
+                    NodeState[task.result_type].model_validate_json(node_state)
+                )
+        return matching_nodes
+
+    def get[R: Result](self, *, task: "Task[Any, R]", index: int = 0) -> NodeState[R]:
+        if (
+            existing_node_state := self._node_states.get((task.name, index))
+        ) is not None:
+            return NodeState[task.result_type].model_validate_json(existing_node_state)
+        else:
+            return NodeState[task.result_type](
+                task_name=task.name,
+                index=index or 0,
+                conversation_state=ConversationState[task.result_type](messages=[]),
+                last_input_hash=0,
+                pinned=False,
+            )
+
+    @contextmanager
+    def run(self) -> Iterator[Self]:
+        self._last_requested_indices = {}
+        yield self
+        self._last_requested_indices = {}
+
+    def dump(self) -> SerializedFlowState:
+        return SerializedFlowState(
+            json.dumps(
+                {
+                    ":".join([task_name, str(index)]): value
+                    for (task_name, index), value in self._node_states.items()
+                }
+            )
+        )
+
+    @classmethod
+    def load(cls, state: SerializedFlowState) -> Self:
+        flow_state = cls()
+        raw_node_states = json.loads(state)
+        new_node_states: dict[tuple[str, int], str] = {}
+        for key, node_state in raw_node_states.items():
+            task_name, index = tuple(key.split(":"))
+            new_node_states[(task_name, int(index))] = node_state
+
+        flow_state._node_states = new_node_states
+        return flow_state
+
+
+_current_flow_state: ContextVar[FlowState | None] = ContextVar(
+    "current_flow_state", default=None
+)
+
+
+class Flow[**P]:
+    def __init__(
+        self, fn: Callable[P, Awaitable[None]], /, *, name: str | None = None
+    ) -> None:
+        self._fn = fn
+        self._name = name
+
+    @property
+    def name(self) -> str:
+        return self._name or self._fn.__name__
+
+    @property
+    def state(self) -> FlowState:
+        state = _current_flow_state.get()
+        if state is None:
+            raise Honk("No current flow state")
+        return state
+
+    @contextmanager
+    def run(self, *, state: FlowState | None = None) -> Iterator[FlowState]:
+        if state is None:
+            state = FlowState()
+
+        old_state = _current_flow_state.get()
+        _current_flow_state.set(state)
+        yield state
+        _current_flow_state.set(old_state)
+
+    async def generate(self, *args: P.args, **kwargs: P.kwargs) -> None:
+        with self.state.run():
+            await self._fn(*args, **kwargs)
+
+
+class Task[**P, R: Result]:
+    def __init__(
+        self,
+        generator: Callable[P, Awaitable[R]],
+        /,
+        *,
+        retries: int = 0,
+    ) -> None:
+        self._generator = generator
+        self._adapter: IAdapter[R] | None = None
+        self._retries = retries
+
+    @property
+    def result_type(self) -> type[R]:
+        result_type = self._generator.__annotations__.get("return")
+        if result_type is None:
+            raise Honk(f"Task {self.name} has no return type annotation")
+        return result_type
+
+    @property
+    def name(self) -> str:
+        return self._generator.__name__
+
+    def adapter(self, adapter: IAdapter[R]) -> Self:
+        self._adapter = adapter
+        return self
+
+    async def generate(
+        self, state: NodeState[R], *args: P.args, **kwargs: P.kwargs
+    ) -> R:
+        input_hash = self.__hash_input(*args, **kwargs)
+        if input_hash != state.last_input_hash:
+            result = await self._generator(*args, **kwargs)
+            state.add_result(result=result, new_input_hash=input_hash, overwrite=True)
+            return result
+        else:
+            if not isinstance(state.conversation_state.messages[-1], GooseResponse):
+                raise Honk(
+                    "Conversation must alternate between User and Result messages"
+                )
+            return state.result
+
+    async def adapt(
+        self, *, flow_state: FlowState, user_message: UserMessage, index: int = 0
+    ) -> R:
+        node_state = flow_state.get(task=self, index=index)
+        if self._adapter is None:
+            raise Honk("No adapter provided for Task")
+
+        node_state.add_user_message(message=user_message)
+        result = await self._adapter(conversation_state=node_state.conversation_state)
+        node_state.add_result(result=result)
+        flow_state.add(node_state)
+
+        return result
+
+    async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        flow_state = self.__get_current_flow_state()
+        node_state = flow_state.get_next(task=self)
+        result = await self.generate(node_state, *args, **kwargs)
+        flow_state.add(node_state)
+        return result
+
+    def __hash_input(self, *args: P.args, **kwargs: P.kwargs) -> int:
+        try:
+            to_hash = str(tuple(args) + tuple(kwargs.values()))
+            return hash(to_hash)
+        except TypeError:
+            raise Honk(f"Unhashable argument to task {self.name}: {args} {kwargs}")
+
+    def __get_current_flow_state(self) -> FlowState:
+        state = _current_flow_state.get()
+        if state is None:
+            raise Honk("No current flow state")
+        return state
+
+
+@overload
+def task[**P, R: Result](generator: Callable[P, Awaitable[R]], /) -> Task[P, R]: ...
+@overload
+def task[**P, R: Result](
+    *, retries: int = 0
+) -> Callable[[Callable[P, Awaitable[R]]], Task[P, R]]: ...
+def task[**P, R: Result](
+    generator: Callable[P, Awaitable[R]] | None = None,
+    /,
+    *,
+    retries: int = 0,
+) -> Task[P, R] | Callable[[Callable[P, Awaitable[R]]], Task[P, R]]:
+    if generator is None:
+
+        def decorator(fn: Callable[P, Awaitable[R]]) -> Task[P, R]:
+            return Task(fn, retries=retries)
+
+        return decorator
+
+    return Task(generator, retries=retries)
+
+
+@overload
+def flow[**P](fn: Callable[P, Awaitable[None]], /) -> Flow[P]: ...
+@overload
+def flow[**P](
+    *, name: str | None = None
+) -> Callable[[Callable[P, Awaitable[None]]], Flow[P]]: ...
+def flow[**P](
+    fn: Callable[P, Awaitable[None]] | None = None, /, *, name: str | 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 decorator
+
+    return Flow(fn, name=name)

{goose_py-0.1.4 → goose_py-0.2.1}/pyproject.toml

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

goose_py-0.1.4/goose/agent.py

@@ -1,75 +0,0 @@
-import logging
-from datetime import datetime
-from typing import Any, Callable
-
-from litellm import acompletion
-from pydantic import BaseModel
-
-from goose.types import (
-    AgentResponse,
-    AssistantMessage,
-    GeminiModel,
-    SystemMessage,
-    UserMessage,
-)
-
-
-class Agent:
-    def __init__(
-        self,
-        *,
-        flow_name: str,
-        run_id: str,
-        logger: Callable[[AgentResponse[Any]], None] | None = None,
-    ) -> None:
-        self.flow_name = flow_name
-        self.run_id = run_id
-        self.logger = logger or logging.info
-
-    async def __call__[R: BaseModel](
-        self,
-        *,
-        messages: list[UserMessage | AssistantMessage],
-        model: GeminiModel,
-        response_model: type[R],
-        task_name: str,
-        system: SystemMessage | None = None,
-    ) -> R:
-        start_time = datetime.now()
-        rendered_messages = [message.render() for message in messages]
-        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.")
-
-        parsed_response = response_model.model_validate_json(
-            response.choices[0].message.content
-        )
-        end_time = datetime.now()
-        agent_response = AgentResponse(
-            response=parsed_response,
-            id=self.run_id,
-            flow_name=self.flow_name,
-            task_name=task_name,
-            model=model,
-            system=system,
-            input_messages=messages,
-            input_tokens=response.usage.prompt_tokens,
-            output_tokens=response.usage.completion_tokens,
-            start_time=start_time,
-            end_time=end_time,
-        )
-
-        self.logger(agent_response)
-        return agent_response.response

goose_py-0.1.4/goose/conversation.py

@@ -1,45 +0,0 @@
-from pydantic import BaseModel
-
-from goose.types import UserMessage
-
-
-class ConversationState[R: BaseModel](BaseModel):
-    user_messages: list[UserMessage]
-    results: list[R]
-
-
-class Conversation[R: BaseModel]:
-    def __init__(
-        self,
-        *,
-        user_messages: list[UserMessage] | None = None,
-        results: list[R] | None = None,
-    ) -> None:
-        self.user_messages = user_messages or []
-        self.results = results or []
-
-    @classmethod
-    def load(cls, *, state: ConversationState[R]) -> "Conversation[R]":
-        return cls(user_messages=state.user_messages, results=state.results)
-
-    @property
-    def current_result(self) -> R:
-        if len(self.results) == 0:
-            raise RuntimeError("No results in conversation")
-
-        return self.results[-1]
-
-    def add_message(self, *, message: UserMessage) -> None:
-        self.user_messages.append(message)
-
-    def add_result(self, *, result: R) -> None:
-        self.results.append(result)
-
-    def replace_last_result(self, *, result: R) -> None:
-        if len(self.results) == 0:
-            self.results.append(result)
-        else:
-            self.results[-1] = result
-
-    def dump(self) -> ConversationState[R]:
-        return ConversationState(user_messages=self.user_messages, results=self.results)

goose_py-0.1.4/goose/core.py

@@ -1,294 +0,0 @@
-import asyncio
-import contextvars
-import inspect
-from collections import defaultdict
-from types import TracebackType
-from typing import Any, Awaitable, Callable, Protocol, Self, overload
-
-from graphlib import TopologicalSorter
-from pydantic import BaseModel
-
-from goose.agent import Agent
-from goose.errors import Honk
-from goose.conversation import Conversation, ConversationState
-from goose.types import AgentResponse, UserMessage
-
-
-class NodeState[R: BaseModel](BaseModel):
-    name: str
-    conversation: ConversationState[R]
-
-    @property
-    def result(self) -> R:
-        return self.conversation.results[-1]
-
-
-class FlowState(BaseModel):
-    nodes: list[NodeState[BaseModel]]
-
-
-class NoResult:
-    pass
-
-
-class IRegenerator[R: BaseModel](Protocol):
-    async def __call__(self, *, result: R, conversation: Conversation[R]) -> R: ...
-
-
-class Task[**P, R: BaseModel]:
-    def __init__(
-        self, generator: Callable[P, Awaitable[R]], /, *, retries: int = 0
-    ) -> None:
-        self.retries = retries
-        self._generator = generator
-        self._regenerator: IRegenerator[R] | None = None
-        self._signature = inspect.signature(generator)
-        self.__validate_fn()
-
-    @property
-    def result_type(self) -> type[R]:
-        return_type = self._generator.__annotations__.get("return")
-        if return_type is None:
-            raise Honk("Task must have a return type annotation")
-
-        return return_type
-
-    @property
-    def name(self) -> str:
-        return self._generator.__name__
-
-    def regenerator(self, regenerator: IRegenerator[R], /) -> Self:
-        self._regenerator = regenerator
-        return self
-
-    async def generate(self, *args: P.args, **kwargs: P.kwargs) -> R:
-        return await self._generator(*args, **kwargs)
-
-    async def regenerate(self, *, result: R, conversation: Conversation[R]) -> R:
-        if self._regenerator is None:
-            raise Honk("Task does not have a regenerator implemented")
-
-        return await self._regenerator(result=result, conversation=conversation)
-
-    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> "Node[R]":
-        arguments = self._signature.bind(*args, **kwargs).arguments
-        return Node(task=self, arguments=arguments, result_type=self.result_type)
-
-    def __validate_fn(self) -> None:
-        if any(
-            param.kind == inspect.Parameter.POSITIONAL_ONLY
-            for param in self._signature.parameters.values()
-        ):
-            raise Honk("Positional-only parameters are not supported in Tasks")
-
-
-class Node[R: BaseModel]:
-    def __init__(
-        self,
-        *,
-        task: Task[Any, R],
-        arguments: dict[str, Any],
-        result_type: type[R],
-        conversation: Conversation[R] | None = None,
-    ) -> None:
-        self.task = task
-        self.arguments = arguments
-        self.result_type = result_type
-        self.conversation = conversation or Conversation[R]()
-        self.name = task.name
-
-        self._result: R | NoResult = NoResult()
-        current_flow = Flow.get_current()
-        if current_flow is None:
-            raise Honk("Cannot create a node without an active flow")
-        self.id = current_flow.add_node(node=self)
-
-    @property
-    def has_result(self) -> bool:
-        return not isinstance(self._result, NoResult)
-
-    @property
-    def result(self) -> R:
-        if isinstance(self._result, NoResult):
-            raise Honk("Cannot access result of a node before it has run")
-        return self._result
-
-    async def generate(self) -> None:
-        self._result = await self.task.generate(**self.arguments)
-        self.conversation.replace_last_result(result=self.result)
-
-    async def regenerate(self, *, message: UserMessage) -> None:
-        self.conversation.add_message(message=message)
-        self._result = await self.task.regenerate(
-            result=self.result, conversation=self.conversation
-        )
-        self.conversation.add_result(result=self.result)
-
-    def dump_state(self) -> NodeState[R]:
-        return NodeState(name=self.name, conversation=self.conversation.dump())
-
-    def load_state(self, *, state: NodeState[R]) -> None:
-        self._result = state.result
-        self.conversation = Conversation[self.result_type].load(
-            state=state.conversation
-        )
-
-    def get_inbound_nodes(self) -> list["Node[BaseModel]"]:
-        def __find_nodes(
-            obj: Any, visited: set[int] | None = None
-        ) -> list["Node[BaseModel]"]:
-            if visited is None:
-                visited = set()
-
-            if isinstance(obj, Node):
-                return [obj]
-            elif isinstance(obj, dict):
-                return [
-                    node
-                    for value in obj.values()
-                    for node in __find_nodes(value, visited)
-                ]
-            elif isinstance(obj, list):
-                return [node for item in obj for node in __find_nodes(item, visited)]
-            elif isinstance(obj, tuple):
-                return [node for item in obj for node in __find_nodes(item, visited)]
-            elif isinstance(obj, set):
-                return [node for item in obj for node in __find_nodes(item, visited)]
-            elif hasattr(obj, "__dict__"):
-                return [
-                    node
-                    for value in obj.__dict__.values()
-                    for node in __find_nodes(value, visited)
-                ]
-            return []
-
-        return __find_nodes(self.arguments)
-
-    def __hash__(self) -> int:
-        return hash(self.id)
-
-
-class Flow:
-    _current: contextvars.ContextVar["Flow | None"] = contextvars.ContextVar(
-        "current_flow", default=None
-    )
-
-    def __init__(
-        self,
-        *,
-        name: str,
-        run_id: str,
-        agent_logger: Callable[[AgentResponse[Any]], None] | None = None,
-    ) -> None:
-        self.name = name
-        self._nodes: list[Node[BaseModel]] = []
-        self._agent = Agent(flow_name=self.name, run_id=run_id, logger=agent_logger)
-
-    @property
-    def agent(self) -> Agent:
-        return self._agent
-
-    def dump_state(self) -> FlowState:
-        return FlowState(nodes=[node.dump_state() for node in self._nodes])
-
-    def load_state(self, *, flow_state: FlowState) -> None:
-        nodes_by_name = {node.name: node for node in self._nodes}
-        for node_state in flow_state.nodes:
-            matching_node = nodes_by_name.get(node_state.name)
-            if matching_node is None:
-                raise Honk(f"Node {node_state.name} from state not found in flow")
-
-            matching_node.load_state(state=node_state)
-
-    async def generate(self) -> None:
-        graph = {node: node.get_inbound_nodes() for node in self._nodes}
-        sorter = TopologicalSorter(graph)
-        sorter.prepare()
-
-        async with asyncio.TaskGroup() as task_group:
-            while sorter.is_active():
-                ready_nodes = list(sorter.get_ready())
-                if ready_nodes:
-                    for node in ready_nodes:
-                        task_group.create_task(node.generate())
-                    sorter.done(*ready_nodes)
-                else:
-                    await asyncio.sleep(0)
-
-    async def regenerate(self, *, target: Node[Any], message: UserMessage) -> None:
-        if not target.has_result:
-            raise Honk("Cannot regenerate a node without a result")
-
-        await target.regenerate(message=message)
-
-        # regenerate all downstream nodes
-        full_graph = {node: node.get_inbound_nodes() for node in self._nodes}
-        reversed_graph: dict[Node[BaseModel], set[Node[BaseModel]]] = defaultdict(set)
-        for node, inbound_nodes in full_graph.items():
-            for inbound_node in inbound_nodes:
-                reversed_graph[inbound_node].add(node)
-
-        subgraph: dict[Node[BaseModel], set[Node[BaseModel]]] = defaultdict(set)
-        queue: list[Node[BaseModel]] = [target]
-
-        while len(queue) > 0:
-            node = queue.pop(0)
-            outbound_nodes = reversed_graph[node]
-            for outbound_node in outbound_nodes:
-                subgraph[outbound_node].add(node)
-                if outbound_node not in subgraph:
-                    queue.append(outbound_node)
-
-        if len(subgraph) > 0:
-            sorter = TopologicalSorter(subgraph)
-            sorter.prepare()
-
-            async with asyncio.TaskGroup() as task_group:
-                while sorter.is_active():
-                    ready_nodes = list(sorter.get_ready())
-                    if len(ready_nodes) > 0:
-                        for node in ready_nodes:
-                            if node != target:
-                                task_group.create_task(node.generate())
-                        sorter.done(*ready_nodes)
-                    else:
-                        await asyncio.sleep(0)
-
-    @classmethod
-    def get_current(cls) -> "Flow | None":
-        return cls._current.get()
-
-    def add_node(self, *, node: Node[Any]) -> str:
-        existing_names = [node.name for node in self._nodes]
-        number = sum(1 for name in existing_names if name == node.name)
-        self._nodes.append(node)
-        node_id = f"{node.name}_{number}"
-        return node_id
-
-    def __enter__(self) -> Self:
-        if self._current.get() is not None:
-            raise Honk("Cannot enter a new flow while another flow is already active")
-        self._current.set(self)
-        return self
-
-    def __exit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_value: BaseException | None,
-        traceback: TracebackType | None,
-    ) -> None:
-        self._current.set(None)
-
-
-@overload
-def task[**P, R: BaseModel](fn: Callable[P, Awaitable[R]], /) -> Task[P, R]: ...
-@overload
-def task[**P, R: BaseModel](
-    *, retries: int = 0
-) -> Callable[[Callable[P, Awaitable[R]]], Task[P, R]]: ...
-def task[**P, R: BaseModel](
-    fn: Callable[P, Awaitable[R]] | None = None, /, *, retries: int = 0
-) -> Task[P, R] | Callable[[Callable[P, Awaitable[R]]], Task[P, R]]:
-    if fn is None:
-        return lambda fn: Task(fn, retries=retries)
-    return Task(fn, retries=retries)