goose-py 0.9.14__tar.gz → 0.9.16__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.9.14 → goose_py-0.9.16}/.github/workflows/publish.yml +0 -1
- {goose_py-0.9.14 → goose_py-0.9.16}/.stubs/litellm/__init__.pyi +3 -7
- {goose_py-0.9.14 → goose_py-0.9.16}/PKG-INFO +1 -1
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/_internal/agent.py +11 -16
- goose_py-0.9.16/goose/_internal/conversation.py +52 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/_internal/state.py +25 -19
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/_internal/task.py +36 -35
- {goose_py-0.9.14 → goose_py-0.9.16}/pyproject.toml +1 -1
- {goose_py-0.9.14 → goose_py-0.9.16}/tests/test_agent.py +2 -3
- goose_py-0.9.16/tests/test_ask.py +124 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/tests/test_refining.py +7 -17
- {goose_py-0.9.14 → goose_py-0.9.16}/tests/test_state.py +5 -6
- {goose_py-0.9.14 → goose_py-0.9.16}/uv.lock +1 -1
- goose_py-0.9.14/goose/_internal/conversation.py +0 -35
- {goose_py-0.9.14 → goose_py-0.9.16}/.envrc +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/.gitignore +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/.python-version +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/.stubs/jsonpath_ng/__init__.pyi +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/Makefile +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/README.md +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/__init__.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/_internal/flow.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/_internal/result.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/_internal/store.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/_internal/types/__init__.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/_internal/types/agent.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/agent.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/errors.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/flow.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/py.typed +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/runs.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/goose/task.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/tests/__init__.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/tests/test_downstream_task.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/tests/test_hashing.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/tests/test_looping.py +0 -0
- {goose_py-0.9.14 → goose_py-0.9.16}/tests/test_regenerate.py +0 -0
@@ -1,4 +1,5 @@
|
|
1
|
-
from typing import
|
1
|
+
from typing import Literal, NotRequired, TypedDict
|
2
|
+
from pydantic import BaseModel
|
2
3
|
|
3
4
|
_LiteLLMGeminiModel = Literal[
|
4
5
|
"vertex_ai/gemini-1.5-flash",
|
@@ -28,11 +29,6 @@ class _LiteLLMMessage(TypedDict):
|
|
28
29
|
content: list[_LiteLLMTextMessageContent | _LiteLLMMediaMessageContent]
|
29
30
|
cache_control: NotRequired[_LiteLLMCacheControl]
|
30
31
|
|
31
|
-
class _LiteLLMResponseFormat(TypedDict):
|
32
|
-
type: Literal["json_object"]
|
33
|
-
response_schema: dict[str, Any] # must be a valid JSON schema
|
34
|
-
enforce_validation: NotRequired[bool]
|
35
|
-
|
36
32
|
class _LiteLLMModelResponseChoiceMessage:
|
37
33
|
role: Literal["assistant"]
|
38
34
|
content: str
|
@@ -60,7 +56,7 @@ async def acompletion(
|
|
60
56
|
*,
|
61
57
|
model: _LiteLLMGeminiModel,
|
62
58
|
messages: list[_LiteLLMMessage],
|
63
|
-
response_format:
|
59
|
+
response_format: type[BaseModel] | None = None,
|
64
60
|
max_tokens: int | None = None,
|
65
61
|
temperature: float = 1.0,
|
66
62
|
) -> ModelResponse: ...
|
@@ -1,6 +1,6 @@
|
|
1
1
|
Metadata-Version: 2.4
|
2
2
|
Name: goose-py
|
3
|
-
Version: 0.9.
|
3
|
+
Version: 0.9.16
|
4
4
|
Summary: A tool for AI workflows based on human-computer collaboration and structured output.
|
5
5
|
Author-email: Nash Taylor <nash@chelle.ai>, Joshua Cook <joshua@chelle.ai>, Michael Sankur <michael@chelle.ai>
|
6
6
|
Requires-Python: >=3.12
|
@@ -7,7 +7,7 @@ from litellm import acompletion
|
|
7
7
|
from pydantic import BaseModel, computed_field
|
8
8
|
|
9
9
|
from .result import Result, TextResult
|
10
|
-
from .types.agent import AIModel,
|
10
|
+
from .types.agent import AIModel, LLMMessage
|
11
11
|
|
12
12
|
|
13
13
|
class AgentResponseDump(TypedDict):
|
@@ -51,8 +51,8 @@ class AgentResponse[R: BaseModel | str](BaseModel):
|
|
51
51
|
flow_name: str
|
52
52
|
task_name: str
|
53
53
|
model: AIModel
|
54
|
-
system:
|
55
|
-
input_messages: list[
|
54
|
+
system: LLMMessage | None = None
|
55
|
+
input_messages: list[LLMMessage]
|
56
56
|
input_tokens: int
|
57
57
|
output_tokens: int
|
58
58
|
start_time: datetime
|
@@ -82,13 +82,13 @@ class AgentResponse[R: BaseModel | str](BaseModel):
|
|
82
82
|
if self.system is None:
|
83
83
|
minimized_system_message = ""
|
84
84
|
else:
|
85
|
-
minimized_system_message = self.system
|
85
|
+
minimized_system_message = self.system
|
86
86
|
for part in minimized_system_message["content"]:
|
87
87
|
if part["type"] == "image_url":
|
88
88
|
part["image_url"] = "__MEDIA__"
|
89
89
|
minimized_system_message = json.dumps(minimized_system_message)
|
90
90
|
|
91
|
-
minimized_input_messages = [message
|
91
|
+
minimized_input_messages = [message for message in self.input_messages]
|
92
92
|
for message in minimized_input_messages:
|
93
93
|
for part in message["content"]:
|
94
94
|
if part["type"] == "image_url":
|
@@ -135,29 +135,24 @@ class Agent:
|
|
135
135
|
async def __call__[R: Result](
|
136
136
|
self,
|
137
137
|
*,
|
138
|
-
messages: list[
|
138
|
+
messages: list[LLMMessage],
|
139
139
|
model: AIModel,
|
140
140
|
task_name: str,
|
141
141
|
response_model: type[R] = TextResult,
|
142
|
-
system:
|
142
|
+
system: LLMMessage | None = None,
|
143
143
|
) -> R:
|
144
144
|
start_time = datetime.now()
|
145
|
-
rendered_messages = [message.render() for message in messages]
|
146
145
|
if system is not None:
|
147
|
-
|
146
|
+
messages.insert(0, system)
|
148
147
|
|
149
148
|
if response_model is TextResult:
|
150
|
-
response = await acompletion(model=model.value, messages=
|
149
|
+
response = await acompletion(model=model.value, messages=messages)
|
151
150
|
parsed_response = response_model.model_validate({"text": response.choices[0].message.content})
|
152
151
|
else:
|
153
152
|
response = await acompletion(
|
154
153
|
model=model.value,
|
155
|
-
messages=
|
156
|
-
response_format=
|
157
|
-
"type": "json_object",
|
158
|
-
"response_schema": response_model.model_json_schema(),
|
159
|
-
"enforce_validation": True,
|
160
|
-
},
|
154
|
+
messages=messages,
|
155
|
+
response_format=response_model,
|
161
156
|
)
|
162
157
|
parsed_response = response_model.model_validate_json(response.choices[0].message.content)
|
163
158
|
|
@@ -0,0 +1,52 @@
|
|
1
|
+
from typing import Self
|
2
|
+
|
3
|
+
from pydantic import BaseModel
|
4
|
+
|
5
|
+
from goose.errors import Honk
|
6
|
+
|
7
|
+
from .result import Result
|
8
|
+
from .types.agent import AssistantMessage, LLMMessage, SystemMessage, UserMessage
|
9
|
+
|
10
|
+
|
11
|
+
class Conversation[R: Result](BaseModel):
|
12
|
+
user_messages: list[UserMessage]
|
13
|
+
assistant_messages: list[R | str]
|
14
|
+
context: SystemMessage | None = None
|
15
|
+
|
16
|
+
@property
|
17
|
+
def awaiting_response(self) -> bool:
|
18
|
+
return len(self.user_messages) == len(self.assistant_messages)
|
19
|
+
|
20
|
+
def render(self) -> list[LLMMessage]:
|
21
|
+
messages: list[LLMMessage] = []
|
22
|
+
if self.context is not None:
|
23
|
+
messages.append(self.context.render())
|
24
|
+
|
25
|
+
for message_index in range(len(self.user_messages)):
|
26
|
+
message = self.assistant_messages[message_index]
|
27
|
+
if isinstance(message, str):
|
28
|
+
messages.append(AssistantMessage(text=message).render())
|
29
|
+
else:
|
30
|
+
messages.append(AssistantMessage(text=message.model_dump_json()).render())
|
31
|
+
|
32
|
+
messages.append(self.user_messages[message_index].render())
|
33
|
+
|
34
|
+
if len(self.assistant_messages) > len(self.user_messages):
|
35
|
+
message = self.assistant_messages[-1]
|
36
|
+
if isinstance(message, str):
|
37
|
+
messages.append(AssistantMessage(text=message).render())
|
38
|
+
else:
|
39
|
+
messages.append(AssistantMessage(text=message.model_dump_json()).render())
|
40
|
+
|
41
|
+
return messages
|
42
|
+
|
43
|
+
def undo(self) -> Self:
|
44
|
+
if len(self.user_messages) == 0:
|
45
|
+
raise Honk("Cannot undo, no user messages")
|
46
|
+
|
47
|
+
if len(self.assistant_messages) == 0:
|
48
|
+
raise Honk("Cannot undo, no assistant messages")
|
49
|
+
|
50
|
+
self.user_messages.pop()
|
51
|
+
self.assistant_messages.pop()
|
52
|
+
return self
|
@@ -4,15 +4,11 @@ from typing import TYPE_CHECKING, Any, NewType, Self
|
|
4
4
|
|
5
5
|
from pydantic import BaseModel, ConfigDict
|
6
6
|
|
7
|
-
from
|
8
|
-
from .
|
9
|
-
|
10
|
-
|
11
|
-
|
12
|
-
UserMessage,
|
13
|
-
)
|
14
|
-
from .conversation import Conversation
|
15
|
-
from .result import Result
|
7
|
+
from goose._internal.agent import Agent, IAgentLogger
|
8
|
+
from goose._internal.conversation import Conversation
|
9
|
+
from goose._internal.result import Result
|
10
|
+
from goose._internal.types.agent import SystemMessage, UserMessage
|
11
|
+
from goose.errors import Honk
|
16
12
|
|
17
13
|
if TYPE_CHECKING:
|
18
14
|
from goose._internal.task import Task
|
@@ -32,10 +28,11 @@ class NodeState[ResultT: Result](BaseModel):
|
|
32
28
|
|
33
29
|
@property
|
34
30
|
def result(self) -> ResultT:
|
35
|
-
|
36
|
-
|
31
|
+
for message in reversed(self.conversation.assistant_messages):
|
32
|
+
if isinstance(message, Result):
|
33
|
+
return message
|
37
34
|
|
38
|
-
|
35
|
+
raise Honk("Node awaiting response, has no result")
|
39
36
|
|
40
37
|
def set_context(self, *, context: SystemMessage) -> Self:
|
41
38
|
self.conversation.context = context
|
@@ -48,24 +45,33 @@ class NodeState[ResultT: Result](BaseModel):
|
|
48
45
|
new_hash: int | None = None,
|
49
46
|
overwrite: bool = False,
|
50
47
|
) -> Self:
|
51
|
-
if overwrite and len(self.conversation.
|
52
|
-
self.conversation.
|
48
|
+
if overwrite and len(self.conversation.assistant_messages) > 0:
|
49
|
+
self.conversation.assistant_messages[-1] = result
|
53
50
|
else:
|
54
|
-
self.conversation.
|
51
|
+
self.conversation.assistant_messages.append(result)
|
55
52
|
if new_hash is not None:
|
56
53
|
self.last_hash = new_hash
|
57
54
|
return self
|
58
55
|
|
56
|
+
def add_answer(self, *, answer: str) -> Self:
|
57
|
+
self.conversation.assistant_messages.append(answer)
|
58
|
+
return self
|
59
|
+
|
59
60
|
def add_user_message(self, *, message: UserMessage) -> Self:
|
60
61
|
self.conversation.user_messages.append(message)
|
61
62
|
return self
|
62
63
|
|
63
64
|
def edit_last_result(self, *, result: ResultT) -> Self:
|
64
|
-
if len(self.conversation.
|
65
|
+
if len(self.conversation.assistant_messages) == 0:
|
65
66
|
raise Honk("Node awaiting response, has no result")
|
66
67
|
|
67
|
-
self.conversation.
|
68
|
-
|
68
|
+
for message_index, message in enumerate(reversed(self.conversation.assistant_messages)):
|
69
|
+
if isinstance(message, Result):
|
70
|
+
index = len(self.conversation.assistant_messages) - message_index - 1
|
71
|
+
self.conversation.assistant_messages[index] = result
|
72
|
+
return self
|
73
|
+
|
74
|
+
raise Honk("Node awaiting response, has no result")
|
69
75
|
|
70
76
|
def undo(self) -> Self:
|
71
77
|
self.conversation.undo()
|
@@ -117,7 +123,7 @@ class FlowRun[FlowArgumentsT: FlowArguments]:
|
|
117
123
|
return NodeState[task.result_type](
|
118
124
|
task_name=task.name,
|
119
125
|
index=index,
|
120
|
-
conversation=Conversation[task.result_type](user_messages=[],
|
126
|
+
conversation=Conversation[task.result_type](user_messages=[], assistant_messages=[]),
|
121
127
|
last_hash=0,
|
122
128
|
)
|
123
129
|
|
@@ -5,11 +5,10 @@ from typing import Any, overload
|
|
5
5
|
from pydantic import BaseModel
|
6
6
|
|
7
7
|
from ..errors import Honk
|
8
|
-
from .agent import Agent, AIModel
|
9
|
-
from .
|
10
|
-
from .result import Result, TextResult
|
8
|
+
from .agent import Agent, AIModel
|
9
|
+
from .result import Result
|
11
10
|
from .state import FlowRun, NodeState, get_current_flow_run
|
12
|
-
from .types.agent import
|
11
|
+
from .types.agent import SystemMessage, UserMessage
|
13
12
|
|
14
13
|
|
15
14
|
class Task[**P, R: Result]:
|
@@ -19,12 +18,11 @@ class Task[**P, R: Result]:
|
|
19
18
|
/,
|
20
19
|
*,
|
21
20
|
retries: int = 0,
|
22
|
-
|
21
|
+
refinement_model: AIModel = AIModel.GEMINI_FLASH,
|
23
22
|
) -> None:
|
24
23
|
self._generator = generator
|
25
24
|
self._retries = retries
|
26
|
-
self.
|
27
|
-
self._adapter_model = adapter_model
|
25
|
+
self._refinement_model = refinement_model
|
28
26
|
|
29
27
|
@property
|
30
28
|
def result_type(self) -> type[R]:
|
@@ -46,6 +44,25 @@ class Task[**P, R: Result]:
|
|
46
44
|
else:
|
47
45
|
return state.result
|
48
46
|
|
47
|
+
async def ask(self, *, user_message: UserMessage, context: SystemMessage | None = None, index: int = 0) -> str:
|
48
|
+
flow_run = self.__get_current_flow_run()
|
49
|
+
node_state = flow_run.get(task=self, index=index)
|
50
|
+
|
51
|
+
if len(node_state.conversation.assistant_messages) == 0:
|
52
|
+
raise Honk("Cannot ask about a task that has not been initially generated")
|
53
|
+
|
54
|
+
node_state.add_user_message(message=user_message)
|
55
|
+
answer = await flow_run.agent(
|
56
|
+
messages=node_state.conversation.render(),
|
57
|
+
model=self._refinement_model,
|
58
|
+
task_name=f"ask--{self.name}",
|
59
|
+
system=context.render() if context is not None else None,
|
60
|
+
)
|
61
|
+
node_state.add_answer(answer=answer.text)
|
62
|
+
flow_run.upsert_node_state(node_state)
|
63
|
+
|
64
|
+
return answer.text
|
65
|
+
|
49
66
|
async def refine(
|
50
67
|
self,
|
51
68
|
*,
|
@@ -56,14 +73,20 @@ class Task[**P, R: Result]:
|
|
56
73
|
flow_run = self.__get_current_flow_run()
|
57
74
|
node_state = flow_run.get(task=self, index=index)
|
58
75
|
|
59
|
-
if len(node_state.conversation.
|
76
|
+
if len(node_state.conversation.assistant_messages) == 0:
|
60
77
|
raise Honk("Cannot refine a task that has not been initially generated")
|
61
78
|
|
62
79
|
if context is not None:
|
63
80
|
node_state.set_context(context=context)
|
64
81
|
node_state.add_user_message(message=user_message)
|
65
82
|
|
66
|
-
result = await
|
83
|
+
result = await flow_run.agent(
|
84
|
+
messages=node_state.conversation.render(),
|
85
|
+
model=self._refinement_model,
|
86
|
+
task_name=f"refine--{self.name}",
|
87
|
+
system=context.render() if context is not None else None,
|
88
|
+
response_model=self.result_type,
|
89
|
+
)
|
67
90
|
node_state.add_result(result=result)
|
68
91
|
flow_run.upsert_node_state(node_state)
|
69
92
|
|
@@ -88,28 +111,6 @@ class Task[**P, R: Result]:
|
|
88
111
|
flow_run.upsert_node_state(node_state)
|
89
112
|
return result
|
90
113
|
|
91
|
-
async def __adapt(self, *, conversation: Conversation[R], agent: Agent) -> R:
|
92
|
-
messages: list[UserMessage | AssistantMessage] = []
|
93
|
-
for message_index in range(len(conversation.user_messages)):
|
94
|
-
user_message = conversation.user_messages[message_index]
|
95
|
-
result = conversation.result_messages[message_index]
|
96
|
-
|
97
|
-
if isinstance(result, TextResult):
|
98
|
-
assistant_text = result.text
|
99
|
-
else:
|
100
|
-
assistant_text = result.model_dump_json()
|
101
|
-
assistant_message = AssistantMessage(text=assistant_text)
|
102
|
-
messages.append(assistant_message)
|
103
|
-
messages.append(user_message)
|
104
|
-
|
105
|
-
return await agent(
|
106
|
-
messages=messages,
|
107
|
-
model=self._adapter_model,
|
108
|
-
task_name=f"adapt--{self.name}",
|
109
|
-
system=conversation.context,
|
110
|
-
response_model=self.result_type,
|
111
|
-
)
|
112
|
-
|
113
114
|
def __hash_task_call(self, *args: P.args, **kwargs: P.kwargs) -> int:
|
114
115
|
def update_hash(argument: Any, current_hash: Any = hashlib.sha256()) -> None:
|
115
116
|
try:
|
@@ -148,20 +149,20 @@ class Task[**P, R: Result]:
|
|
148
149
|
def task[**P, R: Result](generator: Callable[P, Awaitable[R]], /) -> Task[P, R]: ...
|
149
150
|
@overload
|
150
151
|
def task[**P, R: Result](
|
151
|
-
*, retries: int = 0,
|
152
|
+
*, retries: int = 0, refinement_model: AIModel = AIModel.GEMINI_FLASH
|
152
153
|
) -> Callable[[Callable[P, Awaitable[R]]], Task[P, R]]: ...
|
153
154
|
def task[**P, R: Result](
|
154
155
|
generator: Callable[P, Awaitable[R]] | None = None,
|
155
156
|
/,
|
156
157
|
*,
|
157
158
|
retries: int = 0,
|
158
|
-
|
159
|
+
refinement_model: AIModel = AIModel.GEMINI_FLASH,
|
159
160
|
) -> Task[P, R] | Callable[[Callable[P, Awaitable[R]]], Task[P, R]]:
|
160
161
|
if generator is None:
|
161
162
|
|
162
163
|
def decorator(fn: Callable[P, Awaitable[R]]) -> Task[P, R]:
|
163
|
-
return Task(fn, retries=retries,
|
164
|
+
return Task(fn, retries=retries, refinement_model=refinement_model)
|
164
165
|
|
165
166
|
return decorator
|
166
167
|
|
167
|
-
return Task(generator, retries=retries,
|
168
|
+
return Task(generator, retries=retries, refinement_model=refinement_model)
|
@@ -4,8 +4,7 @@ import pytest
|
|
4
4
|
from pytest_mock import MockerFixture
|
5
5
|
|
6
6
|
from goose import Agent, FlowArguments, TextResult, flow, task
|
7
|
-
from goose.
|
8
|
-
from goose.agent import AgentResponse, AIModel, IAgentLogger, UserMessage
|
7
|
+
from goose.agent import AgentResponse, AIModel, IAgentLogger
|
9
8
|
|
10
9
|
|
11
10
|
class TestFlowArguments(FlowArguments):
|
@@ -29,7 +28,7 @@ def mock_litellm(mocker: MockerFixture) -> Mock:
|
|
29
28
|
@task
|
30
29
|
async def use_agent(*, agent: Agent) -> TextResult:
|
31
30
|
return await agent(
|
32
|
-
messages=[
|
31
|
+
messages=[{"role": "user", "content": [{"type": "text", "text": "Hello"}]}],
|
33
32
|
model=AIModel.GEMINI_FLASH_8B,
|
34
33
|
task_name="greet",
|
35
34
|
)
|
@@ -0,0 +1,124 @@
|
|
1
|
+
from unittest.mock import Mock
|
2
|
+
|
3
|
+
import pytest
|
4
|
+
from pytest_mock import MockerFixture
|
5
|
+
|
6
|
+
from goose import Agent, FlowArguments, flow, task
|
7
|
+
from goose._internal.result import TextResult
|
8
|
+
from goose._internal.types.agent import ContentType, MessagePart, UserMessage
|
9
|
+
from goose.errors import Honk
|
10
|
+
|
11
|
+
|
12
|
+
class MockLiteLLMResponse:
|
13
|
+
def __init__(self, *, response: str, prompt_tokens: int, completion_tokens: int) -> None:
|
14
|
+
self.choices = [Mock(message=Mock(content=response))]
|
15
|
+
self.usage = Mock(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens)
|
16
|
+
|
17
|
+
|
18
|
+
@pytest.fixture
|
19
|
+
def mock_litellm(mocker: MockerFixture) -> Mock:
|
20
|
+
return mocker.patch(
|
21
|
+
"goose._internal.agent.acompletion",
|
22
|
+
return_value=MockLiteLLMResponse(response="Here's the explanation!", prompt_tokens=10, completion_tokens=10),
|
23
|
+
)
|
24
|
+
|
25
|
+
|
26
|
+
class MyFlowArguments(FlowArguments):
|
27
|
+
pass
|
28
|
+
|
29
|
+
|
30
|
+
@task
|
31
|
+
async def basic_task(*, flow_arguments: MyFlowArguments) -> TextResult:
|
32
|
+
return TextResult(text="Hello, world!")
|
33
|
+
|
34
|
+
|
35
|
+
@flow
|
36
|
+
async def my_flow(*, flow_arguments: MyFlowArguments, agent: Agent) -> None:
|
37
|
+
await basic_task(flow_arguments=flow_arguments)
|
38
|
+
|
39
|
+
|
40
|
+
@pytest.mark.asyncio
|
41
|
+
@pytest.mark.usefixtures("mock_litellm")
|
42
|
+
async def test_ask_adds_to_conversation():
|
43
|
+
"""Test that ask mode adds messages to conversation but doesn't change result"""
|
44
|
+
|
45
|
+
async with my_flow.start_run(run_id="1") as run:
|
46
|
+
await my_flow.generate(MyFlowArguments())
|
47
|
+
|
48
|
+
# Get the initial result
|
49
|
+
node_state = run.get(task=basic_task)
|
50
|
+
original_result = node_state.result
|
51
|
+
|
52
|
+
# Ask a follow-up question
|
53
|
+
response = await basic_task.ask(
|
54
|
+
user_message=UserMessage(
|
55
|
+
parts=[MessagePart(content="Can you explain how you got that?", content_type=ContentType.TEXT)]
|
56
|
+
)
|
57
|
+
)
|
58
|
+
|
59
|
+
# Verify the response exists and makes sense
|
60
|
+
assert response == "Here's the explanation!"
|
61
|
+
|
62
|
+
# Get updated conversation
|
63
|
+
node_state = run.get(task=basic_task)
|
64
|
+
conversation = node_state.conversation
|
65
|
+
|
66
|
+
# Verify that asking didn't change the original result
|
67
|
+
assert node_state.result == original_result
|
68
|
+
|
69
|
+
# Verify the conversation includes the new messages
|
70
|
+
assert len(conversation.user_messages) == 1
|
71
|
+
assert len(conversation.assistant_messages) == 2 # Original result + response
|
72
|
+
|
73
|
+
# Verify the last messages are our question and the response
|
74
|
+
assert conversation.user_messages[-1].parts[0].content == "Can you explain how you got that?"
|
75
|
+
assert isinstance(conversation.assistant_messages[-1], str)
|
76
|
+
assert conversation.assistant_messages[-1] == "Here's the explanation!"
|
77
|
+
|
78
|
+
|
79
|
+
@pytest.mark.asyncio
|
80
|
+
@pytest.mark.usefixtures("mock_litellm")
|
81
|
+
async def test_ask_requires_completed_task():
|
82
|
+
"""Test that ask mode only works on tasks that haven't been run"""
|
83
|
+
|
84
|
+
async with my_flow.start_run(run_id="2") as run:
|
85
|
+
# Set up flow arguments but don't run the task
|
86
|
+
run.set_flow_arguments(MyFlowArguments())
|
87
|
+
|
88
|
+
# Try to ask before running the task
|
89
|
+
with pytest.raises(Honk, match="Cannot ask about a task that has not been initially generated"):
|
90
|
+
await basic_task.ask(
|
91
|
+
user_message=UserMessage(parts=[MessagePart(content="Can you explain?", content_type=ContentType.TEXT)])
|
92
|
+
)
|
93
|
+
|
94
|
+
|
95
|
+
@pytest.mark.asyncio
|
96
|
+
@pytest.mark.usefixtures("mock_litellm")
|
97
|
+
async def test_ask_multiple_questions():
|
98
|
+
"""Test that we can ask multiple follow-up questions"""
|
99
|
+
|
100
|
+
async with my_flow.start_run(run_id="3") as run:
|
101
|
+
await my_flow.generate(MyFlowArguments())
|
102
|
+
|
103
|
+
# Ask several questions
|
104
|
+
responses: list[str] = []
|
105
|
+
questions = ["Why is that the answer?", "Can you explain it differently?", "What if we added 1 more?"]
|
106
|
+
|
107
|
+
for question in questions:
|
108
|
+
response = await basic_task.ask(
|
109
|
+
user_message=UserMessage(parts=[MessagePart(content=question, content_type=ContentType.TEXT)])
|
110
|
+
)
|
111
|
+
responses.append(response)
|
112
|
+
|
113
|
+
# Verify we got responses for all questions
|
114
|
+
assert len(responses) == len(questions)
|
115
|
+
assert all(response == "Here's the explanation!" for response in responses)
|
116
|
+
|
117
|
+
# Get conversation to verify messages
|
118
|
+
node_state = run.get(task=basic_task)
|
119
|
+
conversation = node_state.conversation
|
120
|
+
|
121
|
+
# Verify the conversation includes all Q&A pairs
|
122
|
+
# Should have: initial result + (question + answer for each question)
|
123
|
+
assert len(conversation.user_messages) == len(questions)
|
124
|
+
assert len(conversation.assistant_messages) == len(questions) + 1 # +1 for initial result
|
@@ -1,9 +1,7 @@
|
|
1
1
|
import random
|
2
2
|
import string
|
3
|
-
from unittest.mock import Mock
|
4
3
|
|
5
4
|
import pytest
|
6
|
-
from pytest_mock import MockerFixture
|
7
5
|
|
8
6
|
from goose import Agent, FlowArguments, Result, flow, task
|
9
7
|
from goose._internal.types.agent import MessagePart
|
@@ -28,15 +26,6 @@ async def generate_random_word(*, n_characters: int) -> GeneratedWord:
|
|
28
26
|
return GeneratedWord(word="".join(random.sample(string.ascii_lowercase, n_characters)))
|
29
27
|
|
30
28
|
|
31
|
-
@pytest.fixture
|
32
|
-
def generate_random_word_adapter(mocker: MockerFixture) -> Mock:
|
33
|
-
return mocker.patch.object(
|
34
|
-
generate_random_word,
|
35
|
-
"_Task__adapt",
|
36
|
-
return_value=GeneratedWord(word="__ADAPTED__"),
|
37
|
-
)
|
38
|
-
|
39
|
-
|
40
29
|
@task
|
41
30
|
async def make_sentence(*, words: list[GeneratedWord]) -> GeneratedSentence:
|
42
31
|
return GeneratedSentence(sentence=" ".join([word.word for word in words]))
|
@@ -49,7 +38,6 @@ async def sentence(*, flow_arguments: MyFlowArguments, agent: Agent) -> None:
|
|
49
38
|
|
50
39
|
|
51
40
|
@pytest.mark.asyncio
|
52
|
-
@pytest.mark.usefixtures("generate_random_word_adapter")
|
53
41
|
async def test_refining() -> None:
|
54
42
|
async with sentence.start_run(run_id="1") as first_run:
|
55
43
|
await sentence.generate(MyFlowArguments())
|
@@ -59,20 +47,22 @@ async def test_refining() -> None:
|
|
59
47
|
|
60
48
|
# imagine this is a new process
|
61
49
|
async with sentence.start_run(run_id="1") as second_run:
|
62
|
-
await generate_random_word.refine(
|
50
|
+
result = await generate_random_word.refine(
|
63
51
|
user_message=UserMessage(parts=[MessagePart(content="Change it")]),
|
64
52
|
context=SystemMessage(parts=[MessagePart(content="Extra info")]),
|
65
53
|
)
|
54
|
+
# Since refine now directly returns the result from the agent call
|
55
|
+
assert isinstance(result, GeneratedWord)
|
66
56
|
|
67
57
|
random_words = second_run.get_all(task=generate_random_word)
|
68
58
|
assert len(random_words) == 3
|
69
|
-
|
70
|
-
assert random_words[
|
71
|
-
assert random_words[
|
59
|
+
# Remove the __ADAPTED__ check since that was specific to the old __adapt method
|
60
|
+
assert isinstance(random_words[0].result, GeneratedWord)
|
61
|
+
assert isinstance(random_words[1].result, GeneratedWord)
|
62
|
+
assert isinstance(random_words[2].result, GeneratedWord)
|
72
63
|
|
73
64
|
|
74
65
|
@pytest.mark.asyncio
|
75
|
-
@pytest.mark.usefixtures("generate_random_word_adapter")
|
76
66
|
async def test_refining_before_generate_fails() -> None:
|
77
67
|
with pytest.raises(Honk):
|
78
68
|
async with sentence.start_run(run_id="2"):
|
@@ -29,11 +29,10 @@ async def generate_random_word(*, n_characters: int) -> GeneratedWord:
|
|
29
29
|
|
30
30
|
@pytest.fixture
|
31
31
|
def generate_random_word_adapter(mocker: MockerFixture) -> Mock:
|
32
|
-
|
33
|
-
|
34
|
-
|
35
|
-
|
36
|
-
)
|
32
|
+
mock_result = GeneratedWord(word="__REFINED__")
|
33
|
+
mock = mocker.patch("goose._internal.agent.Agent.__call__", autospec=True)
|
34
|
+
mock.return_value = mock_result
|
35
|
+
return mock
|
37
36
|
|
38
37
|
|
39
38
|
@task
|
@@ -81,7 +80,7 @@ async def test_state_undo() -> None:
|
|
81
80
|
async with with_state.start_run(run_id="2") as run:
|
82
81
|
generate_random_word.undo()
|
83
82
|
|
84
|
-
assert run.get(task=generate_random_word).result.word != "
|
83
|
+
assert run.get(task=generate_random_word).result.word != "__REFINED__"
|
85
84
|
|
86
85
|
|
87
86
|
@pytest.mark.asyncio
|
@@ -1,35 +0,0 @@
|
|
1
|
-
from typing import Self
|
2
|
-
|
3
|
-
from pydantic import BaseModel
|
4
|
-
|
5
|
-
from .result import Result
|
6
|
-
from .types.agent import AssistantMessage, LLMMessage, SystemMessage, UserMessage
|
7
|
-
|
8
|
-
|
9
|
-
class Conversation[R: Result](BaseModel):
|
10
|
-
user_messages: list[UserMessage]
|
11
|
-
result_messages: list[R]
|
12
|
-
context: SystemMessage | None = None
|
13
|
-
|
14
|
-
@property
|
15
|
-
def awaiting_response(self) -> bool:
|
16
|
-
return len(self.user_messages) == len(self.result_messages)
|
17
|
-
|
18
|
-
def render(self) -> list[LLMMessage]:
|
19
|
-
messages: list[LLMMessage] = []
|
20
|
-
if self.context is not None:
|
21
|
-
messages.append(self.context.render())
|
22
|
-
|
23
|
-
for message_index in range(len(self.user_messages)):
|
24
|
-
messages.append(AssistantMessage(text=self.result_messages[message_index].model_dump_json()).render())
|
25
|
-
messages.append(self.user_messages[message_index].render())
|
26
|
-
|
27
|
-
if len(self.result_messages) > len(self.user_messages):
|
28
|
-
messages.append(AssistantMessage(text=self.result_messages[-1].model_dump_json()).render())
|
29
|
-
|
30
|
-
return messages
|
31
|
-
|
32
|
-
def undo(self) -> Self:
|
33
|
-
self.user_messages.pop()
|
34
|
-
self.result_messages.pop()
|
35
|
-
return self
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|