goose-py 0.6.0__tar.gz → 0.7.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- goose_py-0.7.0/.envrc +1 -0
- goose_py-0.7.0/.github/workflows/publish.yml +75 -0
- goose_py-0.7.0/.gitignore +5 -0
- goose_py-0.7.0/.python-version +1 -0
- goose_py-0.7.0/.stubs/jsonpath_ng/__init__.pyi +14 -0
- goose_py-0.7.0/.stubs/litellm/__init__.pyi +61 -0
- goose_py-0.7.0/PKG-INFO +14 -0
- goose_py-0.7.0/goose/__init__.py +4 -0
- goose_py-0.6.0/goose/_agent.py → goose_py-0.7.0/goose/_internal/agent.py +7 -2
- goose_py-0.6.0/goose/_conversation.py → goose_py-0.7.0/goose/_internal/conversation.py +7 -2
- goose_py-0.6.0/goose/_flow.py → goose_py-0.7.0/goose/_internal/flow.py +5 -5
- goose_py-0.6.0/goose/_state.py → goose_py-0.7.0/goose/_internal/state.py +4 -4
- goose_py-0.6.0/goose/_store.py → goose_py-0.7.0/goose/_internal/store.py +2 -2
- goose_py-0.6.0/goose/_task.py → goose_py-0.7.0/goose/_internal/task.py +5 -5
- goose_py-0.7.0/goose/agent.py +28 -0
- goose_py-0.7.0/goose/flow.py +3 -0
- goose_py-0.7.0/goose/runs.py +4 -0
- goose_py-0.7.0/pyproject.toml +80 -0
- goose_py-0.7.0/tests/__init__.py +0 -0
- goose_py-0.7.0/tests/conftest.py +9 -0
- goose_py-0.7.0/tests/test_agent.py +72 -0
- goose_py-0.7.0/tests/test_complex_flow_arguments.py +23 -0
- goose_py-0.7.0/tests/test_downstream_task.py +37 -0
- goose_py-0.7.0/tests/test_jamming.py +78 -0
- goose_py-0.7.0/tests/test_looping.py +67 -0
- goose_py-0.7.0/tests/test_regenerate.py +48 -0
- goose_py-0.7.0/tests/test_state.py +51 -0
- goose_py-0.7.0/uv.lock +1057 -0
- goose_py-0.6.0/PKG-INFO +0 -31
- goose_py-0.6.0/goose/__init__.py +0 -5
- goose_py-0.6.0/pyproject.toml +0 -66
- {goose_py-0.6.0 → goose_py-0.7.0}/README.md +0 -0
- /goose_py-0.6.0/goose/_result.py → /goose_py-0.7.0/goose/_internal/result.py +0 -0
- {goose_py-0.6.0/goose → goose_py-0.7.0/goose/_internal}/types/__init__.py +0 -0
- {goose_py-0.6.0/goose → goose_py-0.7.0/goose/_internal}/types/agent.py +0 -0
- {goose_py-0.6.0 → goose_py-0.7.0}/goose/errors.py +0 -0
- {goose_py-0.6.0 → goose_py-0.7.0}/goose/py.typed +0 -0
goose_py-0.7.0/.envrc
ADDED
@@ -0,0 +1 @@
|
|
1
|
+
dotenv
|
@@ -0,0 +1,75 @@
|
|
1
|
+
name: CI/CD
|
2
|
+
|
3
|
+
on:
|
4
|
+
push:
|
5
|
+
branches:
|
6
|
+
- main
|
7
|
+
tags:
|
8
|
+
- "v*"
|
9
|
+
pull_request:
|
10
|
+
branches:
|
11
|
+
- main
|
12
|
+
|
13
|
+
jobs:
|
14
|
+
test:
|
15
|
+
runs-on: ubuntu-latest
|
16
|
+
|
17
|
+
steps:
|
18
|
+
- uses: actions/checkout@v4
|
19
|
+
with:
|
20
|
+
fetch-depth: 0
|
21
|
+
|
22
|
+
- name: Set up Python
|
23
|
+
uses: actions/setup-python@v4
|
24
|
+
with:
|
25
|
+
python-version: "3.12"
|
26
|
+
|
27
|
+
- name: Install Poetry
|
28
|
+
uses: snok/install-poetry@v1
|
29
|
+
with:
|
30
|
+
version: 1.8.5
|
31
|
+
virtualenvs-create: true
|
32
|
+
virtualenvs-in-project: true
|
33
|
+
|
34
|
+
- name: Install dependencies
|
35
|
+
run: poetry install --all-extras
|
36
|
+
|
37
|
+
- name: Run tests
|
38
|
+
run: poetry run pytest tests
|
39
|
+
env:
|
40
|
+
GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
|
41
|
+
|
42
|
+
publish:
|
43
|
+
needs: test
|
44
|
+
runs-on: ubuntu-latest
|
45
|
+
if: github.event_name == 'push' && (startsWith(github.ref, 'refs/tags/') || github.ref == 'refs/heads/main')
|
46
|
+
permissions:
|
47
|
+
id-token: write
|
48
|
+
contents: read
|
49
|
+
|
50
|
+
steps:
|
51
|
+
- uses: actions/checkout@v4
|
52
|
+
with:
|
53
|
+
fetch-depth: 0
|
54
|
+
|
55
|
+
- name: Set up Python
|
56
|
+
uses: actions/setup-python@v4
|
57
|
+
with:
|
58
|
+
python-version: "3.12"
|
59
|
+
|
60
|
+
- name: Install Poetry
|
61
|
+
uses: snok/install-poetry@v1
|
62
|
+
with:
|
63
|
+
version: 1.8.5
|
64
|
+
virtualenvs-create: true
|
65
|
+
virtualenvs-in-project: true
|
66
|
+
|
67
|
+
- name: Install dependencies
|
68
|
+
run: poetry install --all-extras
|
69
|
+
|
70
|
+
- name: Build package
|
71
|
+
run: poetry build
|
72
|
+
|
73
|
+
- name: Publish to PyPI
|
74
|
+
if: startsWith(github.ref, 'refs/tags/')
|
75
|
+
uses: pypa/gh-action-pypi-publish@release/v1
|
@@ -0,0 +1 @@
|
|
1
|
+
3.12
|
@@ -0,0 +1,14 @@
|
|
1
|
+
from typing import Any
|
2
|
+
|
3
|
+
class Fields:
|
4
|
+
fields: tuple[str, ...]
|
5
|
+
|
6
|
+
class DatumInContext:
|
7
|
+
value: Any
|
8
|
+
fields: Fields
|
9
|
+
context: "DatumInContext"
|
10
|
+
|
11
|
+
class Jsonpath:
|
12
|
+
def find(self, /, value: Any) -> list[DatumInContext]: ...
|
13
|
+
|
14
|
+
def parse(jsonpath: str, /) -> Jsonpath: ...
|
@@ -0,0 +1,61 @@
|
|
1
|
+
from typing import Any, Literal, NotRequired, TypedDict
|
2
|
+
|
3
|
+
_LiteLLMGeminiModel = Literal[
|
4
|
+
"vertex_ai/gemini-1.5-flash",
|
5
|
+
"vertex_ai/gemini-1.5-pro",
|
6
|
+
"vertex_ai/gemini-1.5-flash-8b",
|
7
|
+
]
|
8
|
+
_MessageRole = Literal["system", "user", "assistant"]
|
9
|
+
|
10
|
+
class _LiteLLMTextMessageContent(TypedDict):
|
11
|
+
type: Literal["text"]
|
12
|
+
text: str
|
13
|
+
|
14
|
+
class _LiteLLMMediaMessageContent(TypedDict):
|
15
|
+
type: Literal["image_url"]
|
16
|
+
image_url: str
|
17
|
+
|
18
|
+
class _LiteLLMCacheControl(TypedDict):
|
19
|
+
type: Literal["ephemeral"]
|
20
|
+
|
21
|
+
class _LiteLLMMessage(TypedDict):
|
22
|
+
role: _MessageRole
|
23
|
+
content: list[_LiteLLMTextMessageContent | _LiteLLMMediaMessageContent]
|
24
|
+
cache_control: NotRequired[_LiteLLMCacheControl]
|
25
|
+
|
26
|
+
class _LiteLLMResponseFormat(TypedDict):
|
27
|
+
type: Literal["json_object"]
|
28
|
+
response_schema: dict[str, Any] # must be a valid JSON schema
|
29
|
+
enforce_validation: NotRequired[bool]
|
30
|
+
|
31
|
+
class _LiteLLMModelResponseChoiceMessage:
|
32
|
+
role: Literal["assistant"]
|
33
|
+
content: str
|
34
|
+
|
35
|
+
class _LiteLLMModelResponseChoice:
|
36
|
+
finish_reason: Literal["stop"]
|
37
|
+
index: int
|
38
|
+
message: _LiteLLMModelResponseChoiceMessage
|
39
|
+
|
40
|
+
class _LiteLLMUsage:
|
41
|
+
completion_tokens: int
|
42
|
+
prompt_tokens: int
|
43
|
+
total_tokens: int
|
44
|
+
|
45
|
+
class ModelResponse:
|
46
|
+
id: str
|
47
|
+
created: int
|
48
|
+
model: _LiteLLMGeminiModel
|
49
|
+
object: Literal["chat.completion"]
|
50
|
+
system_fingerprint: str | None
|
51
|
+
choices: list[_LiteLLMModelResponseChoice]
|
52
|
+
usage: _LiteLLMUsage
|
53
|
+
|
54
|
+
async def acompletion(
|
55
|
+
*,
|
56
|
+
model: _LiteLLMGeminiModel,
|
57
|
+
messages: list[_LiteLLMMessage],
|
58
|
+
response_format: _LiteLLMResponseFormat | None = None,
|
59
|
+
max_tokens: int | None = None,
|
60
|
+
temperature: float = 1.0,
|
61
|
+
) -> ModelResponse: ...
|
goose_py-0.7.0/PKG-INFO
ADDED
@@ -0,0 +1,14 @@
|
|
1
|
+
Metadata-Version: 2.4
|
2
|
+
Name: goose-py
|
3
|
+
Version: 0.7.0
|
4
|
+
Summary: A tool for AI workflows based on human-computer collaboration and structured output.
|
5
|
+
Author-email: Nash Taylor <nash@chelle.ai>, Joshua Cook <joshua@chelle.ai>, Michael Sankur <michael@chelle.ai>
|
6
|
+
Requires-Python: >=3.12
|
7
|
+
Requires-Dist: jsonpath-ng>=1.7.0
|
8
|
+
Requires-Dist: litellm>=1.56.5
|
9
|
+
Requires-Dist: pydantic>=2.8.2
|
10
|
+
Description-Content-Type: text/markdown
|
11
|
+
|
12
|
+
# Goose
|
13
|
+
|
14
|
+
Docs to come.
|
@@ -6,8 +6,13 @@ from typing import Any, ClassVar, Protocol, TypedDict
|
|
6
6
|
from litellm import acompletion
|
7
7
|
from pydantic import BaseModel, computed_field
|
8
8
|
|
9
|
-
from goose.
|
10
|
-
from goose.types.agent import
|
9
|
+
from goose._internal.result import Result, TextResult
|
10
|
+
from goose._internal.types.agent import (
|
11
|
+
AssistantMessage,
|
12
|
+
GeminiModel,
|
13
|
+
SystemMessage,
|
14
|
+
UserMessage,
|
15
|
+
)
|
11
16
|
|
12
17
|
|
13
18
|
class AgentResponseDump(TypedDict):
|
@@ -1,7 +1,12 @@
|
|
1
1
|
from pydantic import BaseModel
|
2
2
|
|
3
|
-
from goose.
|
4
|
-
from goose.types.agent import
|
3
|
+
from goose._internal.result import Result
|
4
|
+
from goose._internal.types.agent import (
|
5
|
+
AssistantMessage,
|
6
|
+
LLMMessage,
|
7
|
+
SystemMessage,
|
8
|
+
UserMessage,
|
9
|
+
)
|
5
10
|
|
6
11
|
|
7
12
|
class Conversation[R: Result](BaseModel):
|
@@ -2,11 +2,11 @@ from contextlib import asynccontextmanager
|
|
2
2
|
from types import CodeType
|
3
3
|
from typing import AsyncIterator, Awaitable, Callable, Protocol, overload
|
4
4
|
|
5
|
-
from goose.
|
6
|
-
from goose.
|
7
|
-
from goose.
|
8
|
-
from goose.
|
9
|
-
from goose.
|
5
|
+
from goose._internal.agent import Agent, IAgentLogger
|
6
|
+
from goose._internal.conversation import Conversation
|
7
|
+
from goose._internal.result import Result
|
8
|
+
from goose._internal.state import FlowRun, get_current_flow_run, set_current_flow_run
|
9
|
+
from goose._internal.store import IFlowRunStore, InMemoryFlowRunStore
|
10
10
|
from goose.errors import Honk
|
11
11
|
|
12
12
|
|
@@ -4,18 +4,18 @@ from typing import TYPE_CHECKING, Any, Self
|
|
4
4
|
|
5
5
|
from pydantic import BaseModel
|
6
6
|
|
7
|
-
from goose.
|
7
|
+
from goose._internal.agent import (
|
8
8
|
Agent,
|
9
9
|
IAgentLogger,
|
10
10
|
SystemMessage,
|
11
11
|
UserMessage,
|
12
12
|
)
|
13
|
-
from goose.
|
14
|
-
from goose.
|
13
|
+
from goose._internal.conversation import Conversation
|
14
|
+
from goose._internal.result import Result
|
15
15
|
from goose.errors import Honk
|
16
16
|
|
17
17
|
if TYPE_CHECKING:
|
18
|
-
from goose.
|
18
|
+
from goose._internal.task import Task
|
19
19
|
|
20
20
|
|
21
21
|
@dataclass
|
@@ -2,8 +2,8 @@ from __future__ import annotations
|
|
2
2
|
|
3
3
|
from typing import Protocol
|
4
4
|
|
5
|
-
from goose.
|
6
|
-
from goose.
|
5
|
+
from goose._internal.flow import FlowRun
|
6
|
+
from goose._internal.state import FlowRunState
|
7
7
|
|
8
8
|
|
9
9
|
class IFlowRunStore(Protocol):
|
@@ -1,11 +1,11 @@
|
|
1
1
|
from typing import Awaitable, Callable, overload
|
2
2
|
|
3
|
-
from goose.
|
4
|
-
from goose.
|
5
|
-
from goose.
|
6
|
-
from goose.
|
3
|
+
from goose._internal.agent import Agent, GeminiModel, SystemMessage, UserMessage
|
4
|
+
from goose._internal.conversation import Conversation
|
5
|
+
from goose._internal.result import Result, TextResult
|
6
|
+
from goose._internal.state import FlowRun, NodeState, get_current_flow_run
|
7
|
+
from goose._internal.types.agent import AssistantMessage
|
7
8
|
from goose.errors import Honk
|
8
|
-
from goose.types.agent import AssistantMessage
|
9
9
|
|
10
10
|
|
11
11
|
class Task[**P, R: Result]:
|
@@ -0,0 +1,28 @@
|
|
1
|
+
from goose._internal.agent import AgentResponse, IAgentLogger
|
2
|
+
from goose._internal.types.agent import (
|
3
|
+
AssistantMessage,
|
4
|
+
GeminiModel,
|
5
|
+
LLMMediaMessagePart,
|
6
|
+
LLMMessage,
|
7
|
+
LLMTextMessagePart,
|
8
|
+
MediaMessagePart,
|
9
|
+
SystemMessage,
|
10
|
+
TextMessagePart,
|
11
|
+
UserMediaContentType,
|
12
|
+
UserMessage,
|
13
|
+
)
|
14
|
+
|
15
|
+
__all__ = [
|
16
|
+
"AgentResponse",
|
17
|
+
"IAgentLogger",
|
18
|
+
"AssistantMessage",
|
19
|
+
"GeminiModel",
|
20
|
+
"LLMMediaMessagePart",
|
21
|
+
"LLMMessage",
|
22
|
+
"LLMTextMessagePart",
|
23
|
+
"MediaMessagePart",
|
24
|
+
"SystemMessage",
|
25
|
+
"TextMessagePart",
|
26
|
+
"UserMediaContentType",
|
27
|
+
"UserMessage",
|
28
|
+
]
|
@@ -0,0 +1,80 @@
|
|
1
|
+
[project]
|
2
|
+
name = "goose-py"
|
3
|
+
version = "0.7.0"
|
4
|
+
description = "A tool for AI workflows based on human-computer collaboration and structured output."
|
5
|
+
readme = "README.md"
|
6
|
+
authors = [
|
7
|
+
{ name = "Nash Taylor", email = "nash@chelle.ai" },
|
8
|
+
{ name = "Joshua Cook", email = "joshua@chelle.ai" },
|
9
|
+
{ name = "Michael Sankur", email = "michael@chelle.ai" },
|
10
|
+
]
|
11
|
+
requires-python = ">=3.12"
|
12
|
+
dependencies = [
|
13
|
+
"jsonpath-ng>=1.7.0",
|
14
|
+
"litellm>=1.56.5",
|
15
|
+
"pydantic>=2.8.2",
|
16
|
+
]
|
17
|
+
|
18
|
+
[build-system]
|
19
|
+
requires = ["hatchling"]
|
20
|
+
build-backend = "hatchling.build"
|
21
|
+
|
22
|
+
[tool.hatch.build.targets.wheel]
|
23
|
+
packages = ["goose"]
|
24
|
+
|
25
|
+
[dependency-groups]
|
26
|
+
dev = [
|
27
|
+
"pyright>=1.1.393",
|
28
|
+
"pytest>=8.3.4",
|
29
|
+
"pytest-asyncio>=0.25.3",
|
30
|
+
"pytest-mock>=3.14.0",
|
31
|
+
"ruff>=0.9.4",
|
32
|
+
]
|
33
|
+
|
34
|
+
|
35
|
+
[tool.ruff]
|
36
|
+
exclude = [
|
37
|
+
"__init__.py",
|
38
|
+
".venv",
|
39
|
+
"**/.venv",
|
40
|
+
"notebooks",
|
41
|
+
]
|
42
|
+
force-exclude = true
|
43
|
+
line-length = 120
|
44
|
+
|
45
|
+
[tool.ruff.lint]
|
46
|
+
select = [ "E", "F", "I", "UP" ]
|
47
|
+
ignore = [ "E501" ]
|
48
|
+
|
49
|
+
|
50
|
+
[tool.ruff.lint.isort]
|
51
|
+
known-first-party = [
|
52
|
+
"goose",
|
53
|
+
]
|
54
|
+
|
55
|
+
[tool.ruff.lint.flake8-tidy-imports]
|
56
|
+
ban-relative-imports = "all"
|
57
|
+
|
58
|
+
|
59
|
+
[tool.pyright]
|
60
|
+
pythonVersion = "3.12"
|
61
|
+
typeCheckingMode = "strict"
|
62
|
+
reportMissingModuleSource = false
|
63
|
+
useLibraryCodeForTypes = false
|
64
|
+
reportImportCycles = true
|
65
|
+
reportUnknownMemberType = false
|
66
|
+
reportUnknownVariableType = false
|
67
|
+
stubPath = ".stubs"
|
68
|
+
venvPath = "."
|
69
|
+
venv = ".venv"
|
70
|
+
|
71
|
+
[tool.pytest.ini_options]
|
72
|
+
filterwarnings = [
|
73
|
+
"ignore::DeprecationWarning",
|
74
|
+
"ignore::SyntaxWarning",
|
75
|
+
"ignore::UserWarning",
|
76
|
+
]
|
77
|
+
addopts = "-v"
|
78
|
+
testpaths = ["tests"]
|
79
|
+
pythonpath = ["."]
|
80
|
+
python_files = "test_*.py"
|
File without changes
|
@@ -0,0 +1,72 @@
|
|
1
|
+
from unittest.mock import Mock
|
2
|
+
|
3
|
+
import pytest
|
4
|
+
from pytest_mock import MockerFixture
|
5
|
+
|
6
|
+
from goose import TextResult, flow, task
|
7
|
+
from goose._internal.agent import Agent, AgentResponse, IAgentLogger
|
8
|
+
from goose.types.agent import GeminiModel, TextMessagePart, UserMessage
|
9
|
+
|
10
|
+
|
11
|
+
class MockLiteLLMResponse:
|
12
|
+
def __init__(
|
13
|
+
self, *, response: str, prompt_tokens: int, completion_tokens: int
|
14
|
+
) -> None:
|
15
|
+
self.choices = [Mock(message=Mock(content=response))]
|
16
|
+
self.usage = Mock(
|
17
|
+
prompt_tokens=prompt_tokens, completion_tokens=completion_tokens
|
18
|
+
)
|
19
|
+
|
20
|
+
|
21
|
+
@pytest.fixture
|
22
|
+
def mock_litellm(mocker: MockerFixture) -> Mock:
|
23
|
+
return mocker.patch(
|
24
|
+
"goose._agent.acompletion",
|
25
|
+
return_value=MockLiteLLMResponse(
|
26
|
+
response="Hello", prompt_tokens=10, completion_tokens=10
|
27
|
+
),
|
28
|
+
)
|
29
|
+
|
30
|
+
|
31
|
+
@task
|
32
|
+
async def use_agent(*, agent: Agent) -> TextResult:
|
33
|
+
return await agent(
|
34
|
+
messages=[UserMessage(parts=[TextMessagePart(text="Hello")])],
|
35
|
+
model=GeminiModel.FLASH_8B,
|
36
|
+
task_name="greet",
|
37
|
+
)
|
38
|
+
|
39
|
+
|
40
|
+
@flow
|
41
|
+
async def agent_flow(*, agent: Agent) -> None:
|
42
|
+
await use_agent(agent=agent)
|
43
|
+
|
44
|
+
|
45
|
+
class CustomLogger(IAgentLogger):
|
46
|
+
logged_responses: list[AgentResponse[TextResult]] = []
|
47
|
+
|
48
|
+
async def __call__(self, *, response: AgentResponse[TextResult]) -> None:
|
49
|
+
self.logged_responses.append(response)
|
50
|
+
|
51
|
+
|
52
|
+
@flow(agent_logger=CustomLogger())
|
53
|
+
async def agent_flow_with_custom_logger(*, agent: Agent) -> None:
|
54
|
+
await use_agent(agent=agent)
|
55
|
+
|
56
|
+
|
57
|
+
@pytest.mark.asyncio
|
58
|
+
@pytest.mark.usefixtures("mock_litellm")
|
59
|
+
async def test_agent() -> None:
|
60
|
+
async with agent_flow.start_run(run_id="1") as run:
|
61
|
+
await agent_flow.generate(agent=run.agent)
|
62
|
+
|
63
|
+
assert run.get(task=use_agent).result.text == "Hello"
|
64
|
+
|
65
|
+
|
66
|
+
@pytest.mark.asyncio
|
67
|
+
@pytest.mark.usefixtures("mock_litellm")
|
68
|
+
async def test_agent_custom_logger() -> None:
|
69
|
+
async with agent_flow_with_custom_logger.start_run(run_id="1") as run:
|
70
|
+
await agent_flow_with_custom_logger.generate(agent=run.agent)
|
71
|
+
|
72
|
+
assert len(CustomLogger.logged_responses) == 1
|
@@ -0,0 +1,23 @@
|
|
1
|
+
import pytest
|
2
|
+
from pydantic import BaseModel
|
3
|
+
|
4
|
+
from goose import flow
|
5
|
+
from goose._agent import Agent
|
6
|
+
|
7
|
+
|
8
|
+
class MyMessage(BaseModel):
|
9
|
+
text: str
|
10
|
+
|
11
|
+
|
12
|
+
@flow
|
13
|
+
async def my_flow(*, message: MyMessage, agent: Agent) -> None:
|
14
|
+
pass
|
15
|
+
|
16
|
+
|
17
|
+
@pytest.mark.asyncio
|
18
|
+
async def test_my_flow() -> None:
|
19
|
+
async with my_flow.start_run(run_id="1") as run:
|
20
|
+
await my_flow.generate(message=MyMessage(text="Hello"), agent=run.agent)
|
21
|
+
|
22
|
+
async with my_flow.start_run(run_id="1") as run:
|
23
|
+
await my_flow.generate(message=MyMessage(text="Hello"), agent=run.agent)
|
@@ -0,0 +1,37 @@
|
|
1
|
+
import random
|
2
|
+
import string
|
3
|
+
|
4
|
+
import pytest
|
5
|
+
|
6
|
+
from goose import Result, flow, task
|
7
|
+
|
8
|
+
|
9
|
+
class GeneratedWord(Result):
|
10
|
+
word: str
|
11
|
+
|
12
|
+
|
13
|
+
@task
|
14
|
+
async def generate_random_word(*, n_characters: int) -> GeneratedWord:
|
15
|
+
return GeneratedWord(
|
16
|
+
word="".join(random.sample(string.ascii_lowercase, n_characters))
|
17
|
+
)
|
18
|
+
|
19
|
+
|
20
|
+
@task
|
21
|
+
async def duplicate_word(*, word: str, times: int) -> GeneratedWord:
|
22
|
+
return GeneratedWord(word="".join([word] * times))
|
23
|
+
|
24
|
+
|
25
|
+
@flow
|
26
|
+
async def downstream_task() -> None:
|
27
|
+
word = await generate_random_word(n_characters=10)
|
28
|
+
await duplicate_word(word=word.word, times=10)
|
29
|
+
|
30
|
+
|
31
|
+
@pytest.mark.asyncio
|
32
|
+
async def test_downstream_task() -> None:
|
33
|
+
async with downstream_task.start_run(run_id="1") as run:
|
34
|
+
await downstream_task.generate()
|
35
|
+
|
36
|
+
duplicated_word = run.get(task=duplicate_word)
|
37
|
+
assert len(duplicated_word.result.word) == 100
|
@@ -0,0 +1,78 @@
|
|
1
|
+
import random
|
2
|
+
import string
|
3
|
+
from unittest.mock import Mock
|
4
|
+
|
5
|
+
import pytest
|
6
|
+
from pytest_mock import MockerFixture
|
7
|
+
|
8
|
+
from goose import Result, flow, task
|
9
|
+
from goose.types.agent import SystemMessage, TextMessagePart, UserMessage
|
10
|
+
|
11
|
+
|
12
|
+
class GeneratedWord(Result):
|
13
|
+
word: str
|
14
|
+
|
15
|
+
|
16
|
+
class GeneratedSentence(Result):
|
17
|
+
sentence: str
|
18
|
+
|
19
|
+
|
20
|
+
@task
|
21
|
+
async def generate_random_word(*, n_characters: int) -> GeneratedWord:
|
22
|
+
return GeneratedWord(
|
23
|
+
word="".join(random.sample(string.ascii_lowercase, n_characters))
|
24
|
+
)
|
25
|
+
|
26
|
+
|
27
|
+
@pytest.fixture
|
28
|
+
def generate_random_word_adapter(mocker: MockerFixture) -> Mock:
|
29
|
+
return mocker.patch.object(
|
30
|
+
generate_random_word,
|
31
|
+
"_Task__adapt",
|
32
|
+
return_value=GeneratedWord(word="__ADAPTED__"),
|
33
|
+
)
|
34
|
+
|
35
|
+
|
36
|
+
@task
|
37
|
+
async def make_sentence(*, words: list[GeneratedWord]) -> GeneratedSentence:
|
38
|
+
return GeneratedSentence(sentence=" ".join([word.word for word in words]))
|
39
|
+
|
40
|
+
|
41
|
+
@flow
|
42
|
+
async def sentence() -> None:
|
43
|
+
words = [await generate_random_word(n_characters=10) for _ in range(3)]
|
44
|
+
await make_sentence(words=words)
|
45
|
+
|
46
|
+
|
47
|
+
@pytest.mark.asyncio
|
48
|
+
@pytest.mark.usefixtures("generate_random_word_adapter")
|
49
|
+
async def test_jamming() -> None:
|
50
|
+
async with sentence.start_run(run_id="1") as first_run:
|
51
|
+
await sentence.generate()
|
52
|
+
|
53
|
+
initial_random_words = first_run.get_all(task=generate_random_word)
|
54
|
+
assert len(initial_random_words) == 3
|
55
|
+
|
56
|
+
# imagine this is a new process
|
57
|
+
async with sentence.start_run(run_id="1") as second_run:
|
58
|
+
await generate_random_word.jam(
|
59
|
+
index=1,
|
60
|
+
user_message=UserMessage(parts=[TextMessagePart(text="Change it")]),
|
61
|
+
context=SystemMessage(parts=[TextMessagePart(text="Extra info")]),
|
62
|
+
)
|
63
|
+
|
64
|
+
random_words = second_run.get_all(task=generate_random_word)
|
65
|
+
assert len(random_words) == 3
|
66
|
+
assert random_words[0].result.word != "__ADAPTED__" # not adapted
|
67
|
+
assert random_words[1].result.word == "__ADAPTED__" # adapted
|
68
|
+
assert random_words[2].result.word != "__ADAPTED__" # not adapted
|
69
|
+
|
70
|
+
# imagine this is a new process
|
71
|
+
async with sentence.start_run(run_id="1") as third_run:
|
72
|
+
await sentence.generate()
|
73
|
+
|
74
|
+
resulting_sentence = third_run.get(task=make_sentence)
|
75
|
+
assert (
|
76
|
+
resulting_sentence.result.sentence
|
77
|
+
== f"{initial_random_words[0].result.word} __ADAPTED__ {initial_random_words[2].result.word}"
|
78
|
+
)
|