goose-py 0.11.17__py3-none-any.whl → 0.11.19__py3-none-any.whl

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/_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/_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.dist-info → goose_py-0.11.19.dist-info}/METADATA

@@ -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.dist-info → goose_py-0.11.19.dist-info}/RECORD

@@ -4,15 +4,15 @@ goose/flow.py,sha256=YsZLBa5I1W27_P6LYGWbtFX8ZYx9vJG3KtENYChHm5E,111
 goose/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 goose/runs.py,sha256=ub-r_gzbUbaIzWXX-jc-dncNxEh6zTfzIkmnDfCSbRI,160
 goose/task.py,sha256=95rspdxETJoY12IHBl3KjnVIdqQnf1jDKlnGWNWOTvQ,53
-goose/_internal/agent.py,sha256=qDPPsOuvTkUpaKXV5aVTIn_yWUs_XZyHDPJaIq49YAI,8704
+goose/_internal/agent.py,sha256=CpVvuyDkNrFLsVIHaf0Tbvp48vgTxyjQdlETsFdbvJ8,9088
 goose/_internal/conversation.py,sha256=vhJwe1pHk2lV60DaB9Tz9KbpzQo7_thRYInPjbIoUTE,1437
 goose/_internal/flow.py,sha256=8MJxlhHYSAzUHZefpF_sRJc37o532OF0X7l3KRopDmc,4115
 goose/_internal/result.py,sha256=vtJMfBxb9skfl8st2tn4hBmEq6qmXiJTme_B5QTgu2M,538
 goose/_internal/state.py,sha256=kA116MpsetsQz6nYodsXOqE3uYz37OTgjC9Vcy_3Qvg,8065
 goose/_internal/store.py,sha256=tWmKfa1-yq1jU6lT3l6kSOmVt2m3H7I1xLMTrxnUDI8,889
-goose/_internal/task.py,sha256=X_eRZxZlf6SwyvF1nIyjoneyqD_TISXqESyxluk63mE,6416
+goose/_internal/task.py,sha256=Qj7Z5CowDAwjcdNJTgfIDEsu7bYr5wUchgJxF4mDp04,6547
 goose/_internal/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose/_internal/types/telemetry.py,sha256=xTJCJo-TW9Pd5QK983HMw9CqJ7vw-BLUkYkbantIP30,4075
-goose_py-0.11.17.dist-info/METADATA,sha256=mid84snXvLt_V2dGQ9Y23BfrTTMuahZaPz3D6T-xsYM,444
-goose_py-0.11.17.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-goose_py-0.11.17.dist-info/RECORD,,
+goose/_internal/types/telemetry.py,sha256=xpfkhx7zCdZjjKU8rPuUJ2UHgqmZ084ZupzbTU36gSM,3791
+goose_py-0.11.19.dist-info/METADATA,sha256=RDrIiV9eVeTK3CEDSBwsQxoubkEbuYN9R03GCqdXK4E,444
+goose_py-0.11.19.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+goose_py-0.11.19.dist-info/RECORD,,