goose-py 0.4.0__tar.gz → 0.4.1__tar.gz

{goose_py-0.4.0 → goose_py-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: goose-py
-Version: 0.4.0
+Version: 0.4.1
 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.0 → goose_py-0.4.1}/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_py-0.4.0 → goose_py-0.4.1}/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,
@@ -25,14 +25,11 @@ from goose.agent import (
 )
 from goose.errors import Honk
 from goose.store import IFlowRunStore, InMemoryFlowRunStore
+from goose.result import Result
 
 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]

goose_py-0.4.1/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 → goose_py-0.4.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "goose-py"
-version = "0.4.0"
+version = "0.4.1"
 description = "A tool for AI workflows based on human-computer collaboration and structured output."
 authors = [
     "Nash Taylor <nash@chelle.ai>",