goose-py 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

goose/agent.py

@@ -7,6 +7,7 @@ from typing import Any, ClassVar, Literal, NotRequired, Protocol, TypedDict
 
 from litellm import acompletion
 from pydantic import BaseModel, computed_field
+from goose.result import Result, TextResult
 
 
 class GeminiModel(StrEnum):
@@ -115,7 +116,7 @@ class AgentResponseDump(TypedDict):
     duration_ms: int
 
 
-class AgentResponse[R: BaseModel](BaseModel):
+class AgentResponse[R: BaseModel | str](BaseModel):
     INPUT_CENTS_PER_MILLION_TOKENS: ClassVar[dict[GeminiModel, float]] = {
         GeminiModel.FLASH_8B: 30,
         GeminiModel.FLASH: 15,
@@ -186,6 +187,12 @@ class AgentResponse[R: BaseModel](BaseModel):
             json.dumps(message) for message in minimized_input_messages
         ]
 
+        output_message = (
+            self.response.model_dump_json()
+            if isinstance(self.response, BaseModel)
+            else self.response
+        )
+
         return {
             "run_id": self.run_id,
             "flow_name": self.flow_name,
@@ -193,7 +200,7 @@ class AgentResponse[R: BaseModel](BaseModel):
             "model": self.model.value,
             "system_message": minimized_system_message,
             "input_messages": minimized_input_messages,
-            "output_message": self.response.model_dump_json(),
+            "output_message": output_message,
             "input_tokens": self.input_tokens,
             "output_tokens": self.output_tokens,
             "input_cost": self.input_cost,
@@ -221,13 +228,13 @@ class Agent:
         self.run_id = run_id
         self.logger = logger
 
-    async def __call__[R: BaseModel](
+    async def __call__[R: Result](
         self,
         *,
         messages: list[UserMessage | AssistantMessage],
         model: GeminiModel,
-        response_model: type[R],
         task_name: str,
+        response_model: type[R] = TextResult,
         system: SystemMessage | None = None,
     ) -> R:
         start_time = datetime.now()
@@ -235,22 +242,25 @@ class Agent:
         if system is not None:
             rendered_messages.insert(0, system.render())
 
-        response = await acompletion(
-            model=model.value,
-            messages=rendered_messages,
-            response_format={
-                "type": "json_object",
-                "response_schema": response_model.model_json_schema(),
-                "enforce_validation": True,
-            },
-        )
-
-        if len(response.choices) == 0:
-            raise RuntimeError("No content returned from LLM call.")
+        if response_model is TextResult:
+            response = await acompletion(model=model.value, messages=rendered_messages)
+            parsed_response = response_model.model_validate(
+                {"text": response.choices[0].message.content}
+            )
+        else:
+            response = await acompletion(
+                model=model.value,
+                messages=rendered_messages,
+                response_format={
+                    "type": "json_object",
+                    "response_schema": response_model.model_json_schema(),
+                    "enforce_validation": True,
+                },
+            )
+            parsed_response = response_model.model_validate_json(
+                response.choices[0].message.content
+            )
 
-        parsed_response = response_model.model_validate_json(
-            response.choices[0].message.content
-        )
         end_time = datetime.now()
         agent_response = AgentResponse(
             response=parsed_response,
@@ -271,4 +281,4 @@ class Agent:
         else:
             logging.info(agent_response.model_dump())
 
-        return agent_response.response
+        return parsed_response

goose/flow.py

@@ -13,7 +13,7 @@ from typing import (
     overload,
 )
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel
 
 from goose.agent import (
     Agent,
@@ -24,15 +24,12 @@ from goose.agent import (
     UserMessage,
 )
 from goose.errors import Honk
+from goose.result import Result
 from goose.store import IFlowRunStore, InMemoryFlowRunStore
 
 SerializedFlowRun = NewType("SerializedFlowRun", str)
 
 
-class Result(BaseModel):
-    model_config = ConfigDict(frozen=True)
-
-
 class Conversation[R: Result](BaseModel):
     user_messages: list[UserMessage]
     result_messages: list[R]
@@ -68,7 +65,9 @@ class Conversation[R: Result](BaseModel):
 class IAdapter[ResultT: Result](Protocol):
     __code__: CodeType
 
-    async def __call__(self, *, conversation: Conversation[ResultT]) -> ResultT: ...
+    async def __call__(
+        self, *, conversation: Conversation[ResultT], agent: Agent
+    ) -> ResultT: ...
 
 
 class NodeState[ResultT: Result](BaseModel):
@@ -357,7 +356,9 @@ class Task[**P, R: Result]:
             node_state.set_context(context=context)
         node_state.add_user_message(message=user_message)
 
-        result = await self._adapter(conversation=node_state.conversation)
+        result = await self._adapter(
+            conversation=node_state.conversation, agent=flow_run.agent
+        )
         node_state.add_result(result=result)
         flow_run.add_node_state(node_state)
 

goose/result.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class Result(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+
+class TextResult(Result):
+    text: str

{goose_py-0.4.0.dist-info → goose_py-0.4.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: goose-py
-Version: 0.4.0
+Version: 0.4.2
 Summary: A tool for AI workflows based on human-computer collaboration and structured output.
 Home-page: https://github.com/chelle-ai/goose
 Keywords: ai,yaml,configuration,llm

goose_py-0.4.2.dist-info/RECORD

@@ -0,0 +1,10 @@
+goose/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+goose/agent.py,sha256=S9nDhK7v4Rsk-ANuq5tIZfZ7TOmCdMEVid2VuK1EWCw,8106
+goose/errors.py,sha256=-0OyZQJWYTRw5YgnCB2_uorVaUsL6Z0QYQO2FqzCiyg,32
+goose/flow.py,sha256=9KF0OKa7k6YIJ7LqI5zpCfaIYoRRPgCQoFeiZ2QU_9M,13500
+goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+goose/result.py,sha256=-eZJn-2sPo7rHZ38Sz6IAHXqiJ-Ss39esEoFGimJEBI,155
+goose/store.py,sha256=4p2BBVAEUS1_Z0iBk5Qk_fPxRQeph64DRzXOFmjIT38,844
+goose_py-0.4.2.dist-info/METADATA,sha256=tSSTAh-EKeEGEyR48MSVVUL--dK1iZ3EAW7G8WD6PlM,1106
+goose_py-0.4.2.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+goose_py-0.4.2.dist-info/RECORD,,

goose_py-0.4.0.dist-info/RECORD

@@ -1,9 +0,0 @@
-goose/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose/agent.py,sha256=pZkKKW2weXEX43p9-pVIpEtwei5nylyQ_YrNZWYFQM0,7687
-goose/errors.py,sha256=-0OyZQJWYTRw5YgnCB2_uorVaUsL6Z0QYQO2FqzCiyg,32
-goose/flow.py,sha256=tzubasKLCOgB_X5-MgC-44OhbwqFZq_7gja3sSElCMI,13478
-goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose/store.py,sha256=4p2BBVAEUS1_Z0iBk5Qk_fPxRQeph64DRzXOFmjIT38,844
-goose_py-0.4.0.dist-info/METADATA,sha256=CJhAf8NiFEOq7pL4Ng4zQbuLMe9Gax-cO1UXJHnYpLY,1106
-goose_py-0.4.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-goose_py-0.4.0.dist-info/RECORD,,