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.
Files changed (37) hide show
  1. goose_py-0.7.0/.envrc +1 -0
  2. goose_py-0.7.0/.github/workflows/publish.yml +75 -0
  3. goose_py-0.7.0/.gitignore +5 -0
  4. goose_py-0.7.0/.python-version +1 -0
  5. goose_py-0.7.0/.stubs/jsonpath_ng/__init__.pyi +14 -0
  6. goose_py-0.7.0/.stubs/litellm/__init__.pyi +61 -0
  7. goose_py-0.7.0/PKG-INFO +14 -0
  8. goose_py-0.7.0/goose/__init__.py +4 -0
  9. goose_py-0.6.0/goose/_agent.py → goose_py-0.7.0/goose/_internal/agent.py +7 -2
  10. goose_py-0.6.0/goose/_conversation.py → goose_py-0.7.0/goose/_internal/conversation.py +7 -2
  11. goose_py-0.6.0/goose/_flow.py → goose_py-0.7.0/goose/_internal/flow.py +5 -5
  12. goose_py-0.6.0/goose/_state.py → goose_py-0.7.0/goose/_internal/state.py +4 -4
  13. goose_py-0.6.0/goose/_store.py → goose_py-0.7.0/goose/_internal/store.py +2 -2
  14. goose_py-0.6.0/goose/_task.py → goose_py-0.7.0/goose/_internal/task.py +5 -5
  15. goose_py-0.7.0/goose/agent.py +28 -0
  16. goose_py-0.7.0/goose/flow.py +3 -0
  17. goose_py-0.7.0/goose/runs.py +4 -0
  18. goose_py-0.7.0/pyproject.toml +80 -0
  19. goose_py-0.7.0/tests/__init__.py +0 -0
  20. goose_py-0.7.0/tests/conftest.py +9 -0
  21. goose_py-0.7.0/tests/test_agent.py +72 -0
  22. goose_py-0.7.0/tests/test_complex_flow_arguments.py +23 -0
  23. goose_py-0.7.0/tests/test_downstream_task.py +37 -0
  24. goose_py-0.7.0/tests/test_jamming.py +78 -0
  25. goose_py-0.7.0/tests/test_looping.py +67 -0
  26. goose_py-0.7.0/tests/test_regenerate.py +48 -0
  27. goose_py-0.7.0/tests/test_state.py +51 -0
  28. goose_py-0.7.0/uv.lock +1057 -0
  29. goose_py-0.6.0/PKG-INFO +0 -31
  30. goose_py-0.6.0/goose/__init__.py +0 -5
  31. goose_py-0.6.0/pyproject.toml +0 -66
  32. {goose_py-0.6.0 → goose_py-0.7.0}/README.md +0 -0
  33. /goose_py-0.6.0/goose/_result.py → /goose_py-0.7.0/goose/_internal/result.py +0 -0
  34. {goose_py-0.6.0/goose → goose_py-0.7.0/goose/_internal}/types/__init__.py +0 -0
  35. {goose_py-0.6.0/goose → goose_py-0.7.0/goose/_internal}/types/agent.py +0 -0
  36. {goose_py-0.6.0 → goose_py-0.7.0}/goose/errors.py +0 -0
  37. {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,5 @@
1
+ __pycache__
2
+ .env
3
+ poetry.lock
4
+ notebooks
5
+ dist
@@ -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: ...
@@ -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.
@@ -0,0 +1,4 @@
1
+ from goose._internal.agent import Agent
2
+ from goose._internal.flow import flow
3
+ from goose._internal.result import Result, TextResult
4
+ from goose._internal.task import task
@@ -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._result import Result, TextResult
10
- from goose.types.agent import AssistantMessage, GeminiModel, SystemMessage, UserMessage
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._result import Result
4
- from goose.types.agent import AssistantMessage, LLMMessage, SystemMessage, UserMessage
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._agent import Agent, IAgentLogger
6
- from goose._conversation import Conversation
7
- from goose._result import Result
8
- from goose._state import FlowRun, get_current_flow_run, set_current_flow_run
9
- from goose._store import IFlowRunStore, InMemoryFlowRunStore
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._agent import (
7
+ from goose._internal.agent import (
8
8
  Agent,
9
9
  IAgentLogger,
10
10
  SystemMessage,
11
11
  UserMessage,
12
12
  )
13
- from goose._conversation import Conversation
14
- from goose._result import Result
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._task import Task
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._flow import FlowRun
6
- from goose._state import FlowRunState
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._agent import Agent, GeminiModel, SystemMessage, UserMessage
4
- from goose._conversation import Conversation
5
- from goose._result import Result, TextResult
6
- from goose._state import FlowRun, NodeState, get_current_flow_run
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,3 @@
1
+ from goose._internal.flow import Flow
2
+
3
+ __all__ = ["Flow"]
@@ -0,0 +1,4 @@
1
+ from goose._internal.state import FlowRun, FlowRunState
2
+ from goose._internal.store import IFlowRunStore
3
+
4
+ __all__ = ["FlowRun", "FlowRunState", "IFlowRunStore"]
@@ -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,9 @@
1
+ import asyncio
2
+ import pytest
3
+
4
+
5
+ @pytest.fixture(scope="session")
6
+ def event_loop():
7
+ loop = asyncio.get_event_loop()
8
+ yield loop
9
+ loop.close()
@@ -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
+ )