goose-py 0.5.1__tar.gz → 0.7.0__tar.gz

goose_py-0.7.0/.envrc

@@ -0,0 +1 @@
+dotenv

goose_py-0.7.0/.github/workflows/publish.yml

@@ -0,0 +1,75 @@
+name: CI/CD
+
+on:
+  push:
+    branches:
+      - main
+    tags:
+      - "v*"
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.12"
+
+      - name: Install Poetry
+        uses: snok/install-poetry@v1
+        with:
+          version: 1.8.5
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+
+      - name: Install dependencies
+        run: poetry install --all-extras
+
+      - name: Run tests
+        run: poetry run pytest tests
+        env:
+          GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
+
+  publish:
+    needs: test
+    runs-on: ubuntu-latest
+    if: github.event_name == 'push' && (startsWith(github.ref, 'refs/tags/') || github.ref == 'refs/heads/main')
+    permissions:
+      id-token: write
+      contents: read
+
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.12"
+
+      - name: Install Poetry
+        uses: snok/install-poetry@v1
+        with:
+          version: 1.8.5
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+
+      - name: Install dependencies
+        run: poetry install --all-extras
+
+      - name: Build package
+        run: poetry build
+
+      - name: Publish to PyPI
+        if: startsWith(github.ref, 'refs/tags/')
+        uses: pypa/gh-action-pypi-publish@release/v1

goose_py-0.7.0/.gitignore

@@ -0,0 +1,5 @@
+__pycache__
+.env
+poetry.lock
+notebooks
+dist

goose_py-0.7.0/.python-version

@@ -0,0 +1 @@
+3.12

goose_py-0.7.0/.stubs/jsonpath_ng/__init__.pyi

@@ -0,0 +1,14 @@
+from typing import Any
+
+class Fields:
+    fields: tuple[str, ...]
+
+class DatumInContext:
+    value: Any
+    fields: Fields
+    context: "DatumInContext"
+
+class Jsonpath:
+    def find(self, /, value: Any) -> list[DatumInContext]: ...
+
+def parse(jsonpath: str, /) -> Jsonpath: ...

goose_py-0.7.0/.stubs/litellm/__init__.pyi

@@ -0,0 +1,61 @@
+from typing import Any, Literal, NotRequired, TypedDict
+
+_LiteLLMGeminiModel = Literal[
+    "vertex_ai/gemini-1.5-flash",
+    "vertex_ai/gemini-1.5-pro",
+    "vertex_ai/gemini-1.5-flash-8b",
+]
+_MessageRole = Literal["system", "user", "assistant"]
+
+class _LiteLLMTextMessageContent(TypedDict):
+    type: Literal["text"]
+    text: str
+
+class _LiteLLMMediaMessageContent(TypedDict):
+    type: Literal["image_url"]
+    image_url: str
+
+class _LiteLLMCacheControl(TypedDict):
+    type: Literal["ephemeral"]
+
+class _LiteLLMMessage(TypedDict):
+    role: _MessageRole
+    content: list[_LiteLLMTextMessageContent | _LiteLLMMediaMessageContent]
+    cache_control: NotRequired[_LiteLLMCacheControl]
+
+class _LiteLLMResponseFormat(TypedDict):
+    type: Literal["json_object"]
+    response_schema: dict[str, Any]  # must be a valid JSON schema
+    enforce_validation: NotRequired[bool]
+
+class _LiteLLMModelResponseChoiceMessage:
+    role: Literal["assistant"]
+    content: str
+
+class _LiteLLMModelResponseChoice:
+    finish_reason: Literal["stop"]
+    index: int
+    message: _LiteLLMModelResponseChoiceMessage
+
+class _LiteLLMUsage:
+    completion_tokens: int
+    prompt_tokens: int
+    total_tokens: int
+
+class ModelResponse:
+    id: str
+    created: int
+    model: _LiteLLMGeminiModel
+    object: Literal["chat.completion"]
+    system_fingerprint: str | None
+    choices: list[_LiteLLMModelResponseChoice]
+    usage: _LiteLLMUsage
+
+async def acompletion(
+    *,
+    model: _LiteLLMGeminiModel,
+    messages: list[_LiteLLMMessage],
+    response_format: _LiteLLMResponseFormat | None = None,
+    max_tokens: int | None = None,
+    temperature: float = 1.0,
+) -> ModelResponse: ...

goose_py-0.7.0/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: goose-py
+Version: 0.7.0
+Summary: A tool for AI workflows based on human-computer collaboration and structured output.
+Author-email: Nash Taylor <nash@chelle.ai>, Joshua Cook <joshua@chelle.ai>, Michael Sankur <michael@chelle.ai>
+Requires-Python: >=3.12
+Requires-Dist: jsonpath-ng>=1.7.0
+Requires-Dist: litellm>=1.56.5
+Requires-Dist: pydantic>=2.8.2
+Description-Content-Type: text/markdown
+
+# Goose
+
+Docs to come.

goose_py-0.7.0/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_py-0.5.1/goose → goose_py-0.7.0/goose/_internal}/agent.py

@@ -1,100 +1,18 @@
 import json
 import logging
 from datetime import datetime
-from enum import StrEnum
-from typing import Any, ClassVar, Literal, NotRequired, Protocol, TypedDict
+from typing import Any, ClassVar, Protocol, TypedDict
 
 from litellm import acompletion
 from pydantic import BaseModel, computed_field
-from goose.result import Result, TextResult
 
-
-class GeminiModel(StrEnum):
-    PRO = "vertex_ai/gemini-1.5-pro"
-    FLASH = "vertex_ai/gemini-1.5-flash"
-    FLASH_8B = "vertex_ai/gemini-1.5-flash-8b"
-
-
-class UserMediaContentType(StrEnum):
-    # images
-    JPEG = "image/jpeg"
-    PNG = "image/png"
-    WEBP = "image/webp"
-
-    # audio
-    MP3 = "audio/mp3"
-    WAV = "audio/wav"
-
-    # files
-    PDF = "application/pdf"
-
-
-class LLMTextMessagePart(TypedDict):
-    type: Literal["text"]
-    text: str
-
-
-class LLMMediaMessagePart(TypedDict):
-    type: Literal["image_url"]
-    image_url: str
-
-
-class CacheControl(TypedDict):
-    type: Literal["ephemeral"]
-
-
-class LLMMessage(TypedDict):
-    role: Literal["user", "assistant", "system"]
-    content: list[LLMTextMessagePart | LLMMediaMessagePart]
-    cache_control: NotRequired[CacheControl]
-
-
-class TextMessagePart(BaseModel):
-    text: str
-
-    def render(self) -> LLMTextMessagePart:
-        return {"type": "text", "text": self.text}
-
-
-class MediaMessagePart(BaseModel):
-    content_type: UserMediaContentType
-    content: str
-
-    def render(self) -> LLMMediaMessagePart:
-        return {
-            "type": "image_url",
-            "image_url": f"data:{self.content_type};base64,{self.content}",
-        }
-
-
-class UserMessage(BaseModel):
-    parts: list[TextMessagePart | MediaMessagePart]
-
-    def render(self) -> LLMMessage:
-        content: LLMMessage = {
-            "role": "user",
-            "content": [part.render() for part in self.parts],
-        }
-        if any(isinstance(part, MediaMessagePart) for part in self.parts):
-            content["cache_control"] = {"type": "ephemeral"}
-        return content
-
-
-class AssistantMessage(BaseModel):
-    text: str
-
-    def render(self) -> LLMMessage:
-        return {"role": "assistant", "content": [{"type": "text", "text": self.text}]}
-
-
-class SystemMessage(BaseModel):
-    parts: list[TextMessagePart | MediaMessagePart]
-
-    def render(self) -> LLMMessage:
-        return {
-            "role": "system",
-            "content": [part.render() for part in self.parts],
-        }
+from goose._internal.result import Result, TextResult
+from goose._internal.types.agent import (
+    AssistantMessage,
+    GeminiModel,
+    SystemMessage,
+    UserMessage,
+)
 
 
 class AgentResponseDump(TypedDict):

goose_py-0.7.0/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_py-0.7.0/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_py-0.7.0/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_py-0.5.1/goose → goose_py-0.7.0/goose/_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)