goose-py 0.5.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

goose/__init__.py

@@ -0,0 +1,4 @@
+from goose._internal.agent import Agent
+from goose._internal.flow import flow
+from goose._internal.result import Result, TextResult
+from goose._internal.task import task

goose/_internal/agent.py

@@ -0,0 +1,201 @@
+import json
+import logging
+from datetime import datetime
+from typing import Any, ClassVar, Protocol, TypedDict
+
+from litellm import acompletion
+from pydantic import BaseModel, computed_field
+
+from goose._internal.result import Result, TextResult
+from goose._internal.types.agent import (
+    AssistantMessage,
+    GeminiModel,
+    SystemMessage,
+    UserMessage,
+)
+
+
+class AgentResponseDump(TypedDict):
+    run_id: str
+    flow_name: str
+    task_name: str
+    model: str
+    system_message: str
+    input_messages: list[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 | str](BaseModel):
+    INPUT_CENTS_PER_MILLION_TOKENS: ClassVar[dict[GeminiModel, float]] = {
+        GeminiModel.FLASH_8B: 30,
+        GeminiModel.FLASH: 15,
+        GeminiModel.PRO: 500,
+    }
+    OUTPUT_CENTS_PER_MILLION_TOKENS: ClassVar[dict[GeminiModel, float]] = {
+        GeminiModel.FLASH_8B: 30,
+        GeminiModel.FLASH: 15,
+        GeminiModel.PRO: 500,
+    }
+
+    response: R
+    run_id: str
+    flow_name: str
+    task_name: str
+    model: GeminiModel
+    system: SystemMessage | None = None
+    input_messages: list[UserMessage | AssistantMessage]
+    input_tokens: int
+    output_tokens: int
+    start_time: datetime
+    end_time: datetime
+
+    @computed_field
+    @property
+    def duration_ms(self) -> int:
+        return int((self.end_time - self.start_time).total_seconds() * 1000)
+
+    @computed_field
+    @property
+    def input_cost(self) -> float:
+        return (
+            self.INPUT_CENTS_PER_MILLION_TOKENS[self.model]
+            * self.input_tokens
+            / 1_000_000
+        )
+
+    @computed_field
+    @property
+    def output_cost(self) -> float:
+        return (
+            self.OUTPUT_CENTS_PER_MILLION_TOKENS[self.model]
+            * self.output_tokens
+            / 1_000_000
+        )
+
+    @computed_field
+    @property
+    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.render()
+            for part in minimized_system_message["content"]:
+                if part["type"] == "image_url":
+                    part["image_url"] = "__MEDIA__"
+            minimized_system_message = json.dumps(minimized_system_message)
+
+        minimized_input_messages = [message.render() for message in self.input_messages]
+        for message in minimized_input_messages:
+            for part in message["content"]:
+                if part["type"] == "image_url":
+                    part["image_url"] = "__MEDIA__"
+        minimized_input_messages = [
+            json.dumps(message) for message in minimized_input_messages
+        ]
+
+        output_message = (
+            self.response.model_dump_json()
+            if isinstance(self.response, BaseModel)
+            else self.response
+        )
+
+        return {
+            "run_id": self.run_id,
+            "flow_name": self.flow_name,
+            "task_name": self.task_name,
+            "model": self.model.value,
+            "system_message": minimized_system_message,
+            "input_messages": minimized_input_messages,
+            "output_message": output_message,
+            "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__(
+        self,
+        *,
+        flow_name: str,
+        run_id: str,
+        logger: IAgentLogger | None = None,
+    ) -> None:
+        self.flow_name = flow_name
+        self.run_id = run_id
+        self.logger = logger
+
+    async def __call__[R: Result](
+        self,
+        *,
+        messages: list[UserMessage | AssistantMessage],
+        model: GeminiModel,
+        task_name: str,
+        response_model: type[R] = TextResult,
+        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())
+
+        if response_model is TextResult:
+            response = await acompletion(model=model.value, messages=rendered_messages)
+            parsed_response = response_model.model_validate(
+                {"text": response.choices[0].message.content}
+            )
+        else:
+            response = await acompletion(
+                model=model.value,
+                messages=rendered_messages,
+                response_format={
+                    "type": "json_object",
+                    "response_schema": response_model.model_json_schema(),
+                    "enforce_validation": True,
+                },
+            )
+            parsed_response = response_model.model_validate_json(
+                response.choices[0].message.content
+            )
+
+        end_time = datetime.now()
+        agent_response = AgentResponse(
+            response=parsed_response,
+            run_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(response=agent_response)
+        else:
+            logging.info(agent_response.model_dump())
+
+        return parsed_response

goose/_internal/conversation.py

@@ -0,0 +1,41 @@
+from pydantic import BaseModel
+
+from goose._internal.result import Result
+from goose._internal.types.agent import (
+    AssistantMessage,
+    LLMMessage,
+    SystemMessage,
+    UserMessage,
+)
+
+
+class Conversation[R: Result](BaseModel):
+    user_messages: list[UserMessage]
+    result_messages: list[R]
+    context: SystemMessage | None = None
+
+    @property
+    def awaiting_response(self) -> bool:
+        return len(self.user_messages) == len(self.result_messages)
+
+    def render(self) -> list[LLMMessage]:
+        messages: list[LLMMessage] = []
+        if self.context is not None:
+            messages.append(self.context.render())
+
+        for message_index in range(len(self.user_messages)):
+            messages.append(
+                AssistantMessage(
+                    text=self.result_messages[message_index].model_dump_json()
+                ).render()
+            )
+            messages.append(self.user_messages[message_index].render())
+
+        if len(self.result_messages) > len(self.user_messages):
+            messages.append(
+                AssistantMessage(
+                    text=self.result_messages[-1].model_dump_json()
+                ).render()
+            )
+
+        return messages

goose/_internal/flow.py

@@ -0,0 +1,106 @@
+from contextlib import asynccontextmanager
+from types import CodeType
+from typing import AsyncIterator, Awaitable, Callable, Protocol, overload
+
+from goose._internal.agent import Agent, IAgentLogger
+from goose._internal.conversation import Conversation
+from goose._internal.result import Result
+from goose._internal.state import FlowRun, get_current_flow_run, set_current_flow_run
+from goose._internal.store import IFlowRunStore, InMemoryFlowRunStore
+from goose.errors import Honk
+
+
+class IAdapter[ResultT: Result](Protocol):
+    __code__: CodeType
+
+    async def __call__(
+        self, *, conversation: Conversation[ResultT], agent: Agent
+    ) -> ResultT: ...
+
+
+class Flow[**P]:
+    def __init__(
+        self,
+        fn: Callable[P, Awaitable[None]],
+        /,
+        *,
+        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:
+        return self._name or self._fn.__name__
+
+    @property
+    def current_run(self) -> FlowRun:
+        run = get_current_flow_run()
+        if run is None:
+            raise Honk("No current flow run")
+        return run
+
+    @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 = existing_run
+
+        old_run = get_current_flow_run()
+        set_current_flow_run(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()
+
+        set_current_flow_run(old_run)
+
+    async def generate(self, *args: P.args, **kwargs: P.kwargs) -> None:
+        flow_run = get_current_flow_run()
+        if flow_run is None:
+            raise Honk("No current flow run")
+
+        flow_run.set_flow_inputs(*args, **kwargs)
+        await self._fn(*args, **kwargs)
+
+    async def regenerate(self) -> None:
+        flow_run = get_current_flow_run()
+        if flow_run is None:
+            raise Honk("No current flow run")
+
+        flow_args, flow_kwargs = flow_run.flow_inputs
+        await self._fn(*flow_args, **flow_kwargs)
+
+
+@overload
+def flow[**P](fn: Callable[P, Awaitable[None]], /) -> Flow[P]: ...
+@overload
+def flow[**P](
+    *,
+    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, store=store, agent_logger=agent_logger)
+
+        return decorator
+
+    return Flow(fn, name=name, store=store, agent_logger=agent_logger)

goose/_internal/state.py

@@ -0,0 +1,190 @@
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Self
+
+from pydantic import BaseModel
+
+from goose._internal.agent import (
+    Agent,
+    IAgentLogger,
+    SystemMessage,
+    UserMessage,
+)
+from goose._internal.conversation import Conversation
+from goose._internal.result import Result
+from goose.errors import Honk
+
+if TYPE_CHECKING:
+    from goose._internal.task import Task
+
+
+@dataclass
+class FlowRunState:
+    node_states: dict[tuple[str, int], str]
+    flow_args: tuple[Any, ...]
+    flow_kwargs: dict[str, Any]
+
+
+class NodeState[ResultT: Result](BaseModel):
+    task_name: str
+    index: int
+    conversation: Conversation[ResultT]
+    last_hash: int
+
+    @property
+    def result(self) -> ResultT:
+        if len(self.conversation.result_messages) == 0:
+            raise Honk("Node awaiting response, has no result")
+
+        return self.conversation.result_messages[-1]
+
+    def set_context(self, *, context: SystemMessage) -> Self:
+        self.conversation.context = context
+        return self
+
+    def add_result(
+        self,
+        *,
+        result: ResultT,
+        new_hash: int | None = None,
+        overwrite: bool = False,
+    ) -> Self:
+        if overwrite and len(self.conversation.result_messages) > 0:
+            self.conversation.result_messages[-1] = result
+        else:
+            self.conversation.result_messages.append(result)
+        if new_hash is not None:
+            self.last_hash = new_hash
+        return self
+
+    def add_user_message(self, *, message: UserMessage) -> Self:
+        self.conversation.user_messages.append(message)
+        return self
+
+
+class FlowRun:
+    def __init__(self) -> None:
+        self._node_states: dict[tuple[str, int], str] = {}
+        self._last_requested_indices: dict[str, int] = {}
+        self._flow_name = ""
+        self._id = ""
+        self._agent: Agent | None = None
+        self._flow_args: tuple[Any, ...] | None = None
+        self._flow_kwargs: dict[str, Any] | None = None
+
+    @property
+    def flow_name(self) -> str:
+        return self._flow_name
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    @property
+    def agent(self) -> Agent:
+        if self._agent is None:
+            raise Honk("Agent is only accessible once a run is started")
+        return self._agent
+
+    @property
+    def flow_inputs(self) -> tuple[tuple[Any, ...], dict[str, Any]]:
+        if self._flow_args is None or self._flow_kwargs is None:
+            raise Honk("This Flow run has not been executed before")
+
+        return self._flow_args, self._flow_kwargs
+
+    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 sorted(matching_nodes, key=lambda node: node.index)
+
+    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,
+                conversation=Conversation[task.result_type](
+                    user_messages=[], result_messages=[]
+                ),
+                last_hash=0,
+            )
+
+    def set_flow_inputs(self, *args: Any, **kwargs: Any) -> None:
+        self._flow_args = args
+        self._flow_kwargs = kwargs
+
+    def add_node_state(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 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, logger=agent_logger
+        )
+
+    def end(self) -> None:
+        self._last_requested_indices = {}
+        self._flow_name = ""
+        self._id = ""
+        self._agent = None
+
+    def clear_node(self, *, task: "Task[Any, Result]", index: int) -> None:
+        key = (task.name, index)
+        if key in self._node_states:
+            del self._node_states[key]
+
+    def dump(self) -> FlowRunState:
+        flow_args, flow_kwargs = self.flow_inputs
+
+        return FlowRunState(
+            node_states=self._node_states,
+            flow_args=flow_args,
+            flow_kwargs=flow_kwargs,
+        )
+
+    @classmethod
+    def load(cls, flow_run_state: FlowRunState, /) -> Self:
+        flow_run = cls()
+        flow_run._node_states = flow_run_state.node_states
+        flow_run._flow_args = flow_run_state.flow_args
+        flow_run._flow_kwargs = flow_run_state.flow_kwargs
+
+        return flow_run
+
+
+_current_flow_run: ContextVar[FlowRun | None] = ContextVar(
+    "current_flow_run", default=None
+)
+
+
+def get_current_flow_run() -> FlowRun | None:
+    return _current_flow_run.get()
+
+
+def set_current_flow_run(flow_run: FlowRun | None) -> None:
+    _current_flow_run.set(flow_run)

goose/{store.py → _internal/store.py}

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Protocol
+from typing import Protocol
 
-if TYPE_CHECKING:
-    from goose.flow import FlowRun
+from goose._internal.flow import FlowRun
+from goose._internal.state import FlowRunState
 
 
 class IFlowRunStore(Protocol):
@@ -16,13 +16,15 @@ class IFlowRunStore(Protocol):
 class InMemoryFlowRunStore(IFlowRunStore):
     def __init__(self, *, flow_name: str) -> None:
         self._flow_name = flow_name
-        self._runs: dict[str, FlowRun] = {}
+        self._runs: dict[str, FlowRunState] = {}
 
     async def get(self, *, run_id: str) -> FlowRun | None:
-        return self._runs.get(run_id)
+        state = self._runs.get(run_id)
+        if state is not None:
+            return FlowRun.load(state)
 
     async def save(self, *, run: FlowRun) -> None:
-        self._runs[run.id] = run
+        self._runs[run.id] = run.dump()
 
     async def delete(self, *, run_id: str) -> None:
         self._runs.pop(run_id, None)