goose-py 0.6.0__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.6.0/goose/_agent.py → goose_py-0.7.0/goose/_internal/agent.py

@@ -6,8 +6,13 @@ from typing import Any, ClassVar, Protocol, TypedDict
 from litellm import acompletion
 from pydantic import BaseModel, computed_field
 
-from goose._result import Result, TextResult
-from goose.types.agent import AssistantMessage, GeminiModel, SystemMessage, UserMessage
+from goose._internal.result import Result, TextResult
+from goose._internal.types.agent import (
+    AssistantMessage,
+    GeminiModel,
+    SystemMessage,
+    UserMessage,
+)
 
 
 class AgentResponseDump(TypedDict):

goose_py-0.6.0/goose/_conversation.py → goose_py-0.7.0/goose/_internal/conversation.py

@@ -1,7 +1,12 @@
 from pydantic import BaseModel
 
-from goose._result import Result
-from goose.types.agent import AssistantMessage, LLMMessage, SystemMessage, UserMessage
+from goose._internal.result import Result
+from goose._internal.types.agent import (
+    AssistantMessage,
+    LLMMessage,
+    SystemMessage,
+    UserMessage,
+)
 
 
 class Conversation[R: Result](BaseModel):

goose_py-0.6.0/goose/_flow.py → goose_py-0.7.0/goose/_internal/flow.py

@@ -2,11 +2,11 @@ from contextlib import asynccontextmanager
 from types import CodeType
 from typing import AsyncIterator, Awaitable, Callable, Protocol, overload
 
-from goose._agent import Agent, IAgentLogger
-from goose._conversation import Conversation
-from goose._result import Result
-from goose._state import FlowRun, get_current_flow_run, set_current_flow_run
-from goose._store import IFlowRunStore, InMemoryFlowRunStore
+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
 
 

goose_py-0.6.0/goose/_state.py → goose_py-0.7.0/goose/_internal/state.py

@@ -4,18 +4,18 @@ from typing import TYPE_CHECKING, Any, Self
 
 from pydantic import BaseModel
 
-from goose._agent import (
+from goose._internal.agent import (
     Agent,
     IAgentLogger,
     SystemMessage,
     UserMessage,
 )
-from goose._conversation import Conversation
-from goose._result import Result
+from goose._internal.conversation import Conversation
+from goose._internal.result import Result
 from goose.errors import Honk
 
 if TYPE_CHECKING:
-    from goose._task import Task
+    from goose._internal.task import Task
 
 
 @dataclass

goose_py-0.6.0/goose/_store.py → goose_py-0.7.0/goose/_internal/store.py

@@ -2,8 +2,8 @@ from __future__ import annotations
 
 from typing import Protocol
 
-from goose._flow import FlowRun
-from goose._state import FlowRunState
+from goose._internal.flow import FlowRun
+from goose._internal.state import FlowRunState
 
 
 class IFlowRunStore(Protocol):

goose_py-0.6.0/goose/_task.py → goose_py-0.7.0/goose/_internal/task.py

@@ -1,11 +1,11 @@
 from typing import Awaitable, Callable, overload
 
-from goose._agent import Agent, GeminiModel, SystemMessage, UserMessage
-from goose._conversation import Conversation
-from goose._result import Result, TextResult
-from goose._state import FlowRun, NodeState, get_current_flow_run
+from goose._internal.agent import Agent, GeminiModel, SystemMessage, UserMessage
+from goose._internal.conversation import Conversation
+from goose._internal.result import Result, TextResult
+from goose._internal.state import FlowRun, NodeState, get_current_flow_run
+from goose._internal.types.agent import AssistantMessage
 from goose.errors import Honk
-from goose.types.agent import AssistantMessage
 
 
 class Task[**P, R: Result]:

goose_py-0.7.0/goose/agent.py

@@ -0,0 +1,28 @@
+from goose._internal.agent import AgentResponse, IAgentLogger
+from goose._internal.types.agent import (
+    AssistantMessage,
+    GeminiModel,
+    LLMMediaMessagePart,
+    LLMMessage,
+    LLMTextMessagePart,
+    MediaMessagePart,
+    SystemMessage,
+    TextMessagePart,
+    UserMediaContentType,
+    UserMessage,
+)
+
+__all__ = [
+    "AgentResponse",
+    "IAgentLogger",
+    "AssistantMessage",
+    "GeminiModel",
+    "LLMMediaMessagePart",
+    "LLMMessage",
+    "LLMTextMessagePart",
+    "MediaMessagePart",
+    "SystemMessage",
+    "TextMessagePart",
+    "UserMediaContentType",
+    "UserMessage",
+]

goose_py-0.7.0/goose/flow.py

@@ -0,0 +1,3 @@
+from goose._internal.flow import Flow
+
+__all__ = ["Flow"]

goose_py-0.7.0/goose/runs.py

@@ -0,0 +1,4 @@
+from goose._internal.state import FlowRun, FlowRunState
+from goose._internal.store import IFlowRunStore
+
+__all__ = ["FlowRun", "FlowRunState", "IFlowRunStore"]

goose_py-0.7.0/pyproject.toml

@@ -0,0 +1,80 @@
+[project]
+name = "goose-py"
+version = "0.7.0"
+description = "A tool for AI workflows based on human-computer collaboration and structured output."
+readme = "README.md"
+authors = [
+    { name = "Nash Taylor", email = "nash@chelle.ai" },
+    { name = "Joshua Cook", email = "joshua@chelle.ai" },
+    { name = "Michael Sankur", email = "michael@chelle.ai" },
+]
+requires-python = ">=3.12"
+dependencies = [
+    "jsonpath-ng>=1.7.0",
+    "litellm>=1.56.5",
+    "pydantic>=2.8.2",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["goose"]
+
+[dependency-groups]
+dev = [
+    "pyright>=1.1.393",
+    "pytest>=8.3.4",
+    "pytest-asyncio>=0.25.3",
+    "pytest-mock>=3.14.0",
+    "ruff>=0.9.4",
+]
+
+
+[tool.ruff]
+exclude = [
+    "__init__.py",
+    ".venv",
+    "**/.venv",
+    "notebooks",
+]
+force-exclude = true
+line-length = 120
+
+[tool.ruff.lint]
+select = [ "E", "F", "I", "UP" ]
+ignore = [ "E501" ]
+
+
+[tool.ruff.lint.isort]
+known-first-party = [
+    "goose",
+]
+
+[tool.ruff.lint.flake8-tidy-imports]
+ban-relative-imports = "all"
+
+
+[tool.pyright]
+pythonVersion = "3.12"
+typeCheckingMode = "strict"
+reportMissingModuleSource = false
+useLibraryCodeForTypes = false
+reportImportCycles = true
+reportUnknownMemberType = false
+reportUnknownVariableType = false
+stubPath = ".stubs"
+venvPath = "."
+venv = ".venv"
+
+[tool.pytest.ini_options]
+filterwarnings = [
+    "ignore::DeprecationWarning",
+    "ignore::SyntaxWarning",
+    "ignore::UserWarning",
+]
+addopts = "-v"
+testpaths = ["tests"]
+pythonpath = ["."]
+python_files = "test_*.py"

goose_py-0.7.0/tests/conftest.py

@@ -0,0 +1,9 @@
+import asyncio
+import pytest
+
+
+@pytest.fixture(scope="session")
+def event_loop():
+    loop = asyncio.get_event_loop()
+    yield loop
+    loop.close()

goose_py-0.7.0/tests/test_agent.py

@@ -0,0 +1,72 @@
+from unittest.mock import Mock
+
+import pytest
+from pytest_mock import MockerFixture
+
+from goose import TextResult, flow, task
+from goose._internal.agent import Agent, AgentResponse, IAgentLogger
+from goose.types.agent import GeminiModel, TextMessagePart, UserMessage
+
+
+class MockLiteLLMResponse:
+    def __init__(
+        self, *, response: str, prompt_tokens: int, completion_tokens: int
+    ) -> None:
+        self.choices = [Mock(message=Mock(content=response))]
+        self.usage = Mock(
+            prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
+        )
+
+
+@pytest.fixture
+def mock_litellm(mocker: MockerFixture) -> Mock:
+    return mocker.patch(
+        "goose._agent.acompletion",
+        return_value=MockLiteLLMResponse(
+            response="Hello", prompt_tokens=10, completion_tokens=10
+        ),
+    )
+
+
+@task
+async def use_agent(*, agent: Agent) -> TextResult:
+    return await agent(
+        messages=[UserMessage(parts=[TextMessagePart(text="Hello")])],
+        model=GeminiModel.FLASH_8B,
+        task_name="greet",
+    )
+
+
+@flow
+async def agent_flow(*, agent: Agent) -> None:
+    await use_agent(agent=agent)
+
+
+class CustomLogger(IAgentLogger):
+    logged_responses: list[AgentResponse[TextResult]] = []
+
+    async def __call__(self, *, response: AgentResponse[TextResult]) -> None:
+        self.logged_responses.append(response)
+
+
+@flow(agent_logger=CustomLogger())
+async def agent_flow_with_custom_logger(*, agent: Agent) -> None:
+    await use_agent(agent=agent)
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("mock_litellm")
+async def test_agent() -> None:
+    async with agent_flow.start_run(run_id="1") as run:
+        await agent_flow.generate(agent=run.agent)
+
+    assert run.get(task=use_agent).result.text == "Hello"
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("mock_litellm")
+async def test_agent_custom_logger() -> None:
+    async with agent_flow_with_custom_logger.start_run(run_id="1") as run:
+        await agent_flow_with_custom_logger.generate(agent=run.agent)
+
+    assert len(CustomLogger.logged_responses) == 1

goose_py-0.7.0/tests/test_complex_flow_arguments.py

@@ -0,0 +1,23 @@
+import pytest
+from pydantic import BaseModel
+
+from goose import flow
+from goose._agent import Agent
+
+
+class MyMessage(BaseModel):
+    text: str
+
+
+@flow
+async def my_flow(*, message: MyMessage, agent: Agent) -> None:
+    pass
+
+
+@pytest.mark.asyncio
+async def test_my_flow() -> None:
+    async with my_flow.start_run(run_id="1") as run:
+        await my_flow.generate(message=MyMessage(text="Hello"), agent=run.agent)
+
+    async with my_flow.start_run(run_id="1") as run:
+        await my_flow.generate(message=MyMessage(text="Hello"), agent=run.agent)

goose_py-0.7.0/tests/test_downstream_task.py

@@ -0,0 +1,37 @@
+import random
+import string
+
+import pytest
+
+from goose import Result, flow, task
+
+
+class GeneratedWord(Result):
+    word: str
+
+
+@task
+async def generate_random_word(*, n_characters: int) -> GeneratedWord:
+    return GeneratedWord(
+        word="".join(random.sample(string.ascii_lowercase, n_characters))
+    )
+
+
+@task
+async def duplicate_word(*, word: str, times: int) -> GeneratedWord:
+    return GeneratedWord(word="".join([word] * times))
+
+
+@flow
+async def downstream_task() -> None:
+    word = await generate_random_word(n_characters=10)
+    await duplicate_word(word=word.word, times=10)
+
+
+@pytest.mark.asyncio
+async def test_downstream_task() -> None:
+    async with downstream_task.start_run(run_id="1") as run:
+        await downstream_task.generate()
+
+    duplicated_word = run.get(task=duplicate_word)
+    assert len(duplicated_word.result.word) == 100

goose_py-0.7.0/tests/test_jamming.py

@@ -0,0 +1,78 @@
+import random
+import string
+from unittest.mock import Mock
+
+import pytest
+from pytest_mock import MockerFixture
+
+from goose import Result, flow, task
+from goose.types.agent import SystemMessage, TextMessagePart, UserMessage
+
+
+class GeneratedWord(Result):
+    word: str
+
+
+class GeneratedSentence(Result):
+    sentence: str
+
+
+@task
+async def generate_random_word(*, n_characters: int) -> GeneratedWord:
+    return GeneratedWord(
+        word="".join(random.sample(string.ascii_lowercase, n_characters))
+    )
+
+
+@pytest.fixture
+def generate_random_word_adapter(mocker: MockerFixture) -> Mock:
+    return mocker.patch.object(
+        generate_random_word,
+        "_Task__adapt",
+        return_value=GeneratedWord(word="__ADAPTED__"),
+    )
+
+
+@task
+async def make_sentence(*, words: list[GeneratedWord]) -> GeneratedSentence:
+    return GeneratedSentence(sentence=" ".join([word.word for word in words]))
+
+
+@flow
+async def sentence() -> None:
+    words = [await generate_random_word(n_characters=10) for _ in range(3)]
+    await make_sentence(words=words)
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("generate_random_word_adapter")
+async def test_jamming() -> None:
+    async with sentence.start_run(run_id="1") as first_run:
+        await sentence.generate()
+
+    initial_random_words = first_run.get_all(task=generate_random_word)
+    assert len(initial_random_words) == 3
+
+    # imagine this is a new process
+    async with sentence.start_run(run_id="1") as second_run:
+        await generate_random_word.jam(
+            index=1,
+            user_message=UserMessage(parts=[TextMessagePart(text="Change it")]),
+            context=SystemMessage(parts=[TextMessagePart(text="Extra info")]),
+        )
+
+    random_words = second_run.get_all(task=generate_random_word)
+    assert len(random_words) == 3
+    assert random_words[0].result.word != "__ADAPTED__"  # not adapted
+    assert random_words[1].result.word == "__ADAPTED__"  # adapted
+    assert random_words[2].result.word != "__ADAPTED__"  # not adapted
+
+    # imagine this is a new process
+    async with sentence.start_run(run_id="1") as third_run:
+        await sentence.generate()
+
+    resulting_sentence = third_run.get(task=make_sentence)
+    assert (
+        resulting_sentence.result.sentence
+        == f"{initial_random_words[0].result.word} __ADAPTED__ {initial_random_words[2].result.word}"
+    )