goose-py 0.11.17__tar.gz → 0.11.19__tar.gz

{goose_py-0.11.17 → goose_py-0.11.19}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: goose-py
-Version: 0.11.17
+Version: 0.11.19
 Summary: A tool for AI workflows based on human-computer collaboration and structured output.
 Author-email: Nash Taylor <nash@chelle.ai>, Joshua Cook <joshua@chelle.ai>, Michael Sankur <michael@chelle.ai>
 Requires-Python: >=3.12
-Requires-Dist: aikernel==0.1.37
+Requires-Dist: aikernel==0.1.39
 Requires-Dist: jsonpath-ng>=1.7.0
 Requires-Dist: pydantic>=2.8.2
 Description-Content-Type: text/markdown

{goose_py-0.11.17 → goose_py-0.11.19}/goose/_internal/agent.py

@@ -4,10 +4,11 @@ from typing import Any, Literal, Protocol, overload
 
 from aikernel import (
     LLMAssistantMessage,
-    LLMModel,
+    LLMModelAlias,
     LLMSystemMessage,
     LLMToolMessage,
     LLMUserMessage,
+    Router,
     llm_structured,
     llm_unstructured,
 )
@@ -40,18 +41,21 @@ class Agent:
         self,
         *,
         messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
-        model: LLMModel,
+        model: LLMModelAlias,
         task_name: str,
+        router: Router,
         response_model: type[R] = TextResult,
     ) -> R:
         start_time = datetime.now()
         typed_messages: list[ExpectedMessage] = [*messages]
 
         if response_model is TextResult:
-            response = await llm_unstructured(model=model, messages=typed_messages)
+            response = await llm_unstructured(model=model, messages=typed_messages, router=router)
             parsed_response = response_model.model_validate({"text": response.text})
         else:
-            response = await llm_structured(model=model, messages=typed_messages, response_model=response_model)
+            response = await llm_structured(
+                model=model, messages=typed_messages, response_model=response_model, router=router
+            )
             parsed_response = response.structured_response
 
         end_time = datetime.now()
@@ -88,12 +92,13 @@ class Agent:
         self,
         *,
         messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
-        model: LLMModel,
+        model: LLMModelAlias,
         task_name: str,
+        router: Router,
     ) -> str:
         start_time = datetime.now()
         typed_messages: list[ExpectedMessage] = [*messages]
-        response = await llm_unstructured(model=model, messages=typed_messages)
+        response = await llm_unstructured(model=model, messages=typed_messages, router=router)
         end_time = datetime.now()
 
         if isinstance(messages[0], LLMSystemMessage):
@@ -128,14 +133,15 @@ class Agent:
         self,
         *,
         messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
-        model: LLMModel,
+        model: LLMModelAlias,
+        router: Router,
         task_name: str,
         response_model: type[R],
     ) -> R:
         start_time = datetime.now()
         typed_messages: list[ExpectedMessage] = [*messages]
         find_replace_response = await llm_structured(
-            model=model, messages=typed_messages, response_model=FindReplaceResponse
+            model=model, messages=typed_messages, response_model=FindReplaceResponse, router=router
         )
         parsed_find_replace_response = find_replace_response.structured_response
         end_time = datetime.now()
@@ -179,7 +185,8 @@ class Agent:
         self,
         *,
         messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
-        model: LLMModel,
+        model: LLMModelAlias,
+        router: Router,
         task_name: str,
         mode: Literal["generate"],
         response_model: type[R],
@@ -190,7 +197,8 @@ class Agent:
         self,
         *,
         messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
-        model: LLMModel,
+        model: LLMModelAlias,
+        router: Router,
         task_name: str,
         mode: Literal["ask"],
         response_model: type[R] = TextResult,
@@ -201,7 +209,8 @@ class Agent:
         self,
         *,
         messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
-        model: LLMModel,
+        model: LLMModelAlias,
+        router: Router,
         task_name: str,
         response_model: type[R],
         mode: Literal["refine"],
@@ -212,7 +221,8 @@ class Agent:
         self,
         *,
         messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
-        model: LLMModel,
+        model: LLMModelAlias,
+        router: Router,
         task_name: str,
         response_model: type[R],
     ) -> R: ...
@@ -221,7 +231,8 @@ class Agent:
         self,
         *,
         messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
-        model: LLMModel,
+        model: LLMModelAlias,
+        router: Router,
         task_name: str,
         response_model: type[R] = TextResult,
         mode: Literal["generate", "ask", "refine"] = "generate",
@@ -229,13 +240,13 @@ class Agent:
         match mode:
             case "generate":
                 return await self.generate(
-                    messages=messages, model=model, task_name=task_name, response_model=response_model
+                    messages=messages, model=model, task_name=task_name, router=router, response_model=response_model
                 )
             case "ask":
-                return await self.ask(messages=messages, model=model, task_name=task_name)
+                return await self.ask(messages=messages, model=model, task_name=task_name, router=router)
             case "refine":
                 return await self.refine(
-                    messages=messages, model=model, task_name=task_name, response_model=response_model
+                    messages=messages, model=model, task_name=task_name, router=router, response_model=response_model
                 )
 
     def __apply_find_replace[R: Result](

{goose_py-0.11.17 → goose_py-0.11.19}/goose/_internal/task.py

@@ -2,7 +2,7 @@ import hashlib
 from collections.abc import Awaitable, Callable
 from typing import Any, overload
 
-from aikernel import LLMModel, LLMSystemMessage, LLMUserMessage
+from aikernel import LLMModelAlias, LLMSystemMessage, LLMUserMessage, Router
 from pydantic import BaseModel
 
 from goose._internal.agent import Agent
@@ -18,11 +18,11 @@ class Task[**P, R: Result]:
         /,
         *,
         retries: int = 0,
-        refinement_model: LLMModel = LLMModel.GEMINI_2_0_FLASH,
+        refinement_model: LLMModelAlias = "gemini-2.0-flash-lite",
     ) -> None:
         self._generator = generator
         self._retries = retries
-        self._refinement_model = refinement_model
+        self._refinement_model: LLMModelAlias = refinement_model
 
     @property
     def result_type(self) -> type[R]:
@@ -45,7 +45,7 @@ class Task[**P, R: Result]:
             return self.result_type.model_validate_json(state.raw_result)
 
     async def ask(
-        self, *, user_message: LLMUserMessage, context: LLMSystemMessage | None = None, index: int = 0
+        self, *, user_message: LLMUserMessage, router: Router, context: LLMSystemMessage | None = None, index: int = 0
     ) -> str:
         flow_run = self.__get_current_flow_run()
         node_state = flow_run.get_state(task=self, index=index)
@@ -62,6 +62,7 @@ class Task[**P, R: Result]:
             model=self._refinement_model,
             task_name=f"ask--{self.name}",
             mode="ask",
+            router=router,
         )
         node_state.add_answer(answer=answer)
         flow_run.upsert_node_state(node_state)
@@ -72,6 +73,7 @@ class Task[**P, R: Result]:
         self,
         *,
         user_message: LLMUserMessage,
+        router: Router,
         context: LLMSystemMessage | None = None,
         index: int = 0,
     ) -> R:
@@ -91,6 +93,7 @@ class Task[**P, R: Result]:
             task_name=f"refine--{self.name}",
             response_model=self.result_type,
             mode="refine",
+            router=router,
         )
         node_state.add_result(result=result.model_dump_json())
         flow_run.upsert_node_state(node_state)
@@ -154,14 +157,14 @@ class Task[**P, R: Result]:
 def task[**P, R: Result](generator: Callable[P, Awaitable[R]], /) -> Task[P, R]: ...
 @overload
 def task[**P, R: Result](
-    *, retries: int = 0, refinement_model: LLMModel = LLMModel.GEMINI_2_0_FLASH
+    *, retries: int = 0, refinement_model: LLMModelAlias = "gemini-2.0-flash-lite"
 ) -> Callable[[Callable[P, Awaitable[R]]], Task[P, R]]: ...
 def task[**P, R: Result](
     generator: Callable[P, Awaitable[R]] | None = None,
     /,
     *,
     retries: int = 0,
-    refinement_model: LLMModel = LLMModel.GEMINI_2_0_FLASH,
+    refinement_model: LLMModelAlias = "gemini-2.0-flash-lite",
 ) -> Task[P, R] | Callable[[Callable[P, Awaitable[R]]], Task[P, R]]:
     if generator is None:
 

{goose_py-0.11.17 → goose_py-0.11.19}/goose/_internal/types/telemetry.py

@@ -2,7 +2,7 @@ import json
 from datetime import datetime
 from typing import ClassVar, TypedDict
 
-from aikernel import LiteLLMMessage, LLMModel
+from aikernel import LiteLLMMessage, LLMModelAlias
 from pydantic import BaseModel, computed_field
 
 from goose.errors import Honk
@@ -27,28 +27,24 @@ class AgentResponseDump(TypedDict):
 
 
 class AgentResponse[R: BaseModel | str](BaseModel):
-    INPUT_DOLLARS_PER_MILLION_TOKENS: ClassVar[dict[LLMModel, float]] = {
-        LLMModel.VERTEX_GEMINI_2_0_FLASH: 0.30,
-        LLMModel.VERTEX_GEMINI_2_0_FLASH_LITE: 0.15,
-        LLMModel.VERTEX_GEMINI_2_0_PRO_EXP_02_05: 5.00,
-        LLMModel.GEMINI_2_0_FLASH: 0.30,
-        LLMModel.GEMINI_2_0_FLASH_LITE: 0.15,
-        LLMModel.GEMINI_2_0_PRO_EXP_02_05: 5.00,
+    INPUT_DOLLARS_PER_MILLION_TOKENS: ClassVar[dict[LLMModelAlias, float]] = {
+        "gemini-2.0-flash": 0.30,
+        "gemini-2.0-flash-lite": 0.15,
+        "claude-3.5-sonnet": 3.00,
+        "claude-3.7-sonnet": 3.00,
     }
-    OUTPUT_DOLLARS_PER_MILLION_TOKENS: ClassVar[dict[LLMModel, float]] = {
-        LLMModel.VERTEX_GEMINI_2_0_FLASH: 0.30,
-        LLMModel.VERTEX_GEMINI_2_0_FLASH_LITE: 0.15,
-        LLMModel.VERTEX_GEMINI_2_0_PRO_EXP_02_05: 5.00,
-        LLMModel.GEMINI_2_0_FLASH: 0.30,
-        LLMModel.GEMINI_2_0_FLASH_LITE: 0.15,
-        LLMModel.GEMINI_2_0_PRO_EXP_02_05: 5.00,
+    OUTPUT_DOLLARS_PER_MILLION_TOKENS: ClassVar[dict[LLMModelAlias, float]] = {
+        "gemini-2.0-flash": 0.30,
+        "gemini-2.0-flash-lite": 0.15,
+        "claude-3.5-sonnet": 15.00,
+        "claude-3.7-sonnet": 15.00,
     }
 
     response: R
     run_id: str
     flow_name: str
     task_name: str
-    model: LLMModel
+    model: LLMModelAlias
     system: LiteLLMMessage | None = None
     input_messages: list[LiteLLMMessage]
     input_tokens: int
@@ -102,7 +98,7 @@ class AgentResponse[R: BaseModel | str](BaseModel):
             "run_id": self.run_id,
             "flow_name": self.flow_name,
             "task_name": self.task_name,
-            "model": self.model.value,
+            "model": self.model,
             "system_message": minimized_system_message,
             "input_messages": minimized_input_messages,
             "output_message": output_message,

{goose_py-0.11.17 → goose_py-0.11.19}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "goose-py"
-version = "0.11.17"
+version = "0.11.19"
 description = "A tool for AI workflows based on human-computer collaboration and structured output."
 readme = "README.md"
 authors = [
@@ -11,7 +11,7 @@ authors = [
 requires-python = ">=3.12"
 dependencies = [
     "jsonpath-ng>=1.7.0",
-    "aikernel==0.1.37",
+    "aikernel==0.1.39",
     "pydantic>=2.8.2",
 ]
 

{goose_py-0.11.17 → goose_py-0.11.19}/tests/test_agent.py

@@ -1,7 +1,7 @@
 from unittest.mock import Mock
 
 import pytest
-from aikernel import LLMMessagePart, LLMModel, LLMUserMessage
+from aikernel import LLMMessagePart, LLMUserMessage, Router
 from pytest_mock import MockerFixture
 
 from goose import Agent, AgentResponse, FlowArguments, TextResult, flow, task
@@ -22,12 +22,18 @@ def mock_llm_unstructured(mocker: MockerFixture) -> Mock:
 
 @task
 async def use_agent(*, agent: Agent) -> TextResult:
+    router = Router(
+        model_list=[
+            {"model_name": "gemini-2.0-flash-lite", "litellm_params": {"model": "gemini/gemini-2.0-flash-lite"}}
+        ]
+    )
     return await agent(
         messages=[LLMUserMessage(parts=[LLMMessagePart(content="Hello")])],
-        model=LLMModel.GEMINI_2_0_FLASH_LITE,
+        model="gemini-2.0-flash-lite",
         task_name="greet",
         mode="generate",
         response_model=TextResult,
+        router=router,
     )
 
 

{goose_py-0.11.17 → goose_py-0.11.19}/tests/test_ask.py

@@ -1,13 +1,17 @@
 from unittest.mock import Mock
 
 import pytest
-from aikernel import LLMAssistantMessage, LLMMessagePart, LLMUserMessage
+from aikernel import LLMAssistantMessage, LLMMessagePart, LLMUserMessage, Router
 from pytest_mock import MockerFixture
 
 from goose import Agent, FlowArguments, flow, task
 from goose._internal.result import TextResult
 from goose.errors import Honk
 
+ROUTER = Router(
+    model_list=[{"model_name": "gemini-2.0-flash-lite", "litellm_params": {"model": "gemini/gemini-2.0-flash-lite"}}]
+)
+
 
 class MockLiteLLMResponse:
     def __init__(self, *, response: str, prompt_tokens: int, completion_tokens: int) -> None:
@@ -51,7 +55,8 @@ async def test_ask_adds_to_conversation():
 
         # Ask a follow-up question
         response = await basic_task.ask(
-            user_message=LLMUserMessage(parts=[LLMMessagePart(content="Can you explain how you got that?")])
+            user_message=LLMUserMessage(parts=[LLMMessagePart(content="Can you explain how you got that?")]),
+            router=ROUTER,
         )
 
         # Verify the response exists and makes sense
@@ -85,7 +90,10 @@ async def test_ask_requires_completed_task():
 
         # Try to ask before running the task
         with pytest.raises(Honk, match="Cannot ask about a task that has not been initially generated"):
-            await basic_task.ask(user_message=LLMUserMessage(parts=[LLMMessagePart(content="Can you explain?")]))
+            await basic_task.ask(
+                user_message=LLMUserMessage(parts=[LLMMessagePart(content="Can you explain?")]),
+                router=ROUTER,
+            )
 
 
 @pytest.mark.asyncio
@@ -101,7 +109,10 @@ async def test_ask_multiple_questions():
         questions = ["Why is that the answer?", "Can you explain it differently?", "What if we added 1 more?"]
 
         for question in questions:
-            response = await basic_task.ask(user_message=LLMUserMessage(parts=[LLMMessagePart(content=question)]))
+            response = await basic_task.ask(
+                user_message=LLMUserMessage(parts=[LLMMessagePart(content=question)]),
+                router=ROUTER,
+            )
             responses.append(response)
 
         # Verify we got responses for all questions

{goose_py-0.11.17 → goose_py-0.11.19}/tests/test_refining.py

@@ -3,13 +3,17 @@ import string
 from unittest.mock import Mock
 
 import pytest
-from aikernel import LLMMessagePart, LLMSystemMessage, LLMUserMessage
+from aikernel import LLMMessagePart, LLMSystemMessage, LLMUserMessage, Router
 from pytest_mock import MockerFixture
 
 from goose import Agent, FlowArguments, Result, flow, task
 from goose._internal.result import FindReplaceResponse, Replacement
 from goose.errors import Honk
 
+ROUTER = Router(
+    model_list=[{"model_name": "gemini-2.0-flash-lite", "litellm_params": {"model": "gemini/gemini-2.0-flash-lite"}}]
+)
+
 
 class MyFlowArguments(FlowArguments):
     pass
@@ -63,6 +67,7 @@ async def test_refining() -> None:
             index=0,
             user_message=LLMUserMessage(parts=[LLMMessagePart(content="Change it")]),
             context=LLMSystemMessage(parts=[LLMMessagePart(content="Extra info")]),
+            router=ROUTER,
         )
 
     initial_random_words = first_run.get_all_results(task=generate_random_word)
@@ -73,6 +78,7 @@ async def test_refining() -> None:
         result = await generate_random_word.refine(
             user_message=LLMUserMessage(parts=[LLMMessagePart(content="Change it")]),
             context=LLMSystemMessage(parts=[LLMMessagePart(content="Extra info")]),
+            router=ROUTER,
         )
         # Since refine now directly returns the result from the agent call
         assert isinstance(result, GeneratedWord)
@@ -91,4 +97,5 @@ async def test_refining_before_generate_fails() -> None:
             await generate_random_word.refine(
                 user_message=LLMUserMessage(parts=[LLMMessagePart(content="Change it")]),
                 context=LLMSystemMessage(parts=[LLMMessagePart(content="Extra info")]),
+                router=ROUTER,
             )

{goose_py-0.11.17 → goose_py-0.11.19}/tests/test_state.py

@@ -75,6 +75,7 @@ async def test_state_undo() -> None:
             index=0,
             user_message=LLMUserMessage(parts=[LLMMessagePart(content="Change it")]),
             context=LLMSystemMessage(parts=[LLMMessagePart(content="Extra info")]),
+            router=Mock(),
         )
 
     async with with_state.start_run(run_id="2") as run:

{goose_py-0.11.17 → goose_py-0.11.19}/uv.lock

@@ -3,15 +3,15 @@ requires-python = ">=3.12"
 
 [[package]]
 name = "aikernel"
-version = "0.1.37"
+version = "0.1.39"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "litellm" },
     { name = "pydantic" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/43/c7/b9bc3786867e14cc93cd722ea52835c11348119d109cf0d973883f0214f7/aikernel-0.1.37.tar.gz", hash = "sha256:9e841785d05508962324d355757486eface886327e4abcd95d4ac725866ffa92", size = 71982 }
+sdist = { url = "https://files.pythonhosted.org/packages/ab/7e/1da152279e50e6d651b4f29cc6ee950f4e5aad5c64b8b4889df036c20acd/aikernel-0.1.39.tar.gz", hash = "sha256:13e21eaadd66499d3fb00c82d42e6d1ae4022e69017133d1d04da58636a8e399", size = 72243 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/90/9f/cf7a04bd49e663f6436e47f62e85e6ead3dc28b25da2b4573398fa0849cd/aikernel-0.1.37-py3-none-any.whl", hash = "sha256:b6543af231427ba25ea699ab6646943f5caa607999efc85573fbd100193bab65", size = 9431 },
+    { url = "https://files.pythonhosted.org/packages/2f/86/d55e7c00a1d2df6507ab3b41fb5933be2e24c5ad53c8e2879fff8c2008fc/aikernel-0.1.39-py3-none-any.whl", hash = "sha256:604426ae04880e227cffdef491f3ee305661b957f46153465c8e69740313b368", size = 9972 },
 ]
 
 [[package]]
@@ -345,7 +345,7 @@ wheels = [
 
 [[package]]
 name = "goose-py"
-version = "0.11.17"
+version = "0.11.19"
 source = { editable = "." }
 dependencies = [
     { name = "aikernel" },
@@ -365,7 +365,7 @@ dev = [
 
 [package.metadata]
 requires-dist = [
-    { name = "aikernel", specifier = "==0.1.37" },
+    { name = "aikernel", specifier = "==0.1.39" },
     { name = "jsonpath-ng", specifier = ">=1.7.0" },
     { name = "pydantic", specifier = ">=2.8.2" },
 ]