goose-py 0.10.2__py3-none-any.whl → 0.11.1__py3-none-any.whl

goose/__init__.py

@@ -1,6 +1,7 @@
-from ._internal.agent import Agent
-from ._internal.flow import FlowArguments, flow
-from ._internal.result import Result, TextResult
-from ._internal.task import task
+from goose._internal.agent import Agent
+from goose._internal.flow import FlowArguments, flow
+from goose._internal.result import Result, TextResult
+from goose._internal.task import task
+from goose._internal.types.telemetry import AgentResponse
 
-__all__ = ["Agent", "flow", "FlowArguments", "Result", "TextResult", "task"]
+__all__ = ["Agent", "flow", "FlowArguments", "Result", "TextResult", "task", "AgentResponse"]

goose/_internal/agent.py

@@ -2,11 +2,17 @@ import logging
 from datetime import datetime
 from typing import Any, Literal, Protocol, overload
 
-from litellm import acompletion
+from aikernel import (
+    LLMAssistantMessage,
+    LLMModel,
+    LLMSystemMessage,
+    LLMUserMessage,
+    llm_structured,
+    llm_unstructured,
+)
 from pydantic import ValidationError
 
 from goose._internal.result import FindReplaceResponse, Result, TextResult
-from goose._internal.types.agent import AIModel, AssistantMessage, SystemMessage, UserMessage
 from goose._internal.types.telemetry import AgentResponse
 from goose.errors import Honk
 
@@ -30,42 +36,39 @@ class Agent:
     async def generate[R: Result](
         self,
         *,
-        messages: list[UserMessage | AssistantMessage],
-        model: AIModel,
+        messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
+        model: LLMModel,
         task_name: str,
         response_model: type[R] = TextResult,
-        system: SystemMessage | None = None,
     ) -> R:
-        rendered_messages = [message.render() for message in messages]
-        rendered_system = system.render() if system is not None else None
-
-        completion_messages = (
-            [rendered_system] + rendered_messages if rendered_system is not None else rendered_messages
-        )
-
         start_time = datetime.now()
+
         if response_model is TextResult:
-            response = await acompletion(model=model.value, messages=completion_messages)
-            parsed_response = response_model.model_validate({"text": response.choices[0].message.content})
+            response = await llm_unstructured(model=model, messages=messages)
+            parsed_response = response_model.model_validate({"text": response.text})
         else:
-            response = await acompletion(
-                model=model.value,
-                messages=completion_messages,
-                response_format=response_model,
-            )
-            parsed_response = response_model.model_validate_json(response.choices[0].message.content)
+            response = await llm_structured(model=model, messages=messages, response_model=response_model)
+            parsed_response = response.structured_response
 
         end_time = datetime.now()
+
+        if isinstance(messages[0], LLMSystemMessage):
+            system = messages[0].render()
+            input_messages = [message.render() for message in messages[1:]]
+        else:
+            system = None
+            input_messages = [message.render() for message in messages]
+
         agent_response = AgentResponse(
             response=parsed_response,
             run_id=self.run_id,
             flow_name=self.flow_name,
             task_name=task_name,
             model=model,
-            system=rendered_system,
-            input_messages=rendered_messages,
-            input_tokens=response.usage.prompt_tokens,
-            output_tokens=response.usage.completion_tokens,
+            system=system,
+            input_messages=input_messages,
+            input_tokens=response.usage.input_tokens,
+            output_tokens=response.usage.output_tokens,
             start_time=start_time,
             end_time=end_time,
         )
@@ -80,32 +83,31 @@ class Agent:
     async def ask(
         self,
         *,
-        messages: list[UserMessage | AssistantMessage],
-        model: AIModel,
+        messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
+        model: LLMModel,
         task_name: str,
-        system: SystemMessage | None = None,
     ) -> str:
-        rendered_messages = [message.render() for message in messages]
-        rendered_system = system.render() if system is not None else None
-
-        completion_messages = (
-            [rendered_system] + rendered_messages if rendered_system is not None else rendered_messages
-        )
-
         start_time = datetime.now()
-        response = await acompletion(model=model.value, messages=completion_messages)
-
+        response = await llm_unstructured(model=model, messages=messages)
         end_time = datetime.now()
+
+        if isinstance(messages[0], LLMSystemMessage):
+            system = messages[0].render()
+            input_messages = [message.render() for message in messages[1:]]
+        else:
+            system = None
+            input_messages = [message.render() for message in messages]
+
         agent_response = AgentResponse(
-            response=response.choices[0].message.content,
+            response=response.text,
             run_id=self.run_id,
             flow_name=self.flow_name,
             task_name=task_name,
             model=model,
-            system=rendered_system,
-            input_messages=rendered_messages,
-            input_tokens=response.usage.prompt_tokens,
-            output_tokens=response.usage.completion_tokens,
+            system=system,
+            input_messages=input_messages,
+            input_tokens=response.usage.input_tokens,
+            output_tokens=response.usage.output_tokens,
             start_time=start_time,
             end_time=end_time,
         )
@@ -115,44 +117,38 @@ class Agent:
         else:
             logging.info(agent_response.model_dump())
 
-        return response.choices[0].message.content
+        return response.text
 
     async def refine[R: Result](
         self,
         *,
-        messages: list[UserMessage | AssistantMessage],
-        model: AIModel,
+        messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
+        model: LLMModel,
         task_name: str,
         response_model: type[R],
-        system: SystemMessage | None = None,
     ) -> R:
         start_time = datetime.now()
+        find_replace_response = await llm_structured(model=model, messages=messages, response_model=FindReplaceResponse)
+        parsed_find_replace_response = find_replace_response.structured_response
+        end_time = datetime.now()
 
-        rendered_messages = [message.render() for message in messages]
-        rendered_system = system.render() if system is not None else None
-
-        completion_messages = (
-            [rendered_system] + rendered_messages if rendered_system is not None else rendered_messages
-        )
-
-        find_replace_response = await acompletion(
-            model=model.value, messages=completion_messages, response_format=FindReplaceResponse
-        )
-        parsed_find_replace_response = FindReplaceResponse.model_validate_json(
-            find_replace_response.choices[0].message.content
-        )
+        if isinstance(messages[0], LLMSystemMessage):
+            system = messages[0].render()
+            input_messages = [message.render() for message in messages[1:]]
+        else:
+            system = None
+            input_messages = [message.render() for message in messages]
 
-        end_time = datetime.now()
         agent_response = AgentResponse(
             response=parsed_find_replace_response,
             run_id=self.run_id,
             flow_name=self.flow_name,
             task_name=task_name,
             model=model,
-            system=rendered_system,
-            input_messages=rendered_messages,
-            input_tokens=find_replace_response.usage.prompt_tokens,
-            output_tokens=find_replace_response.usage.completion_tokens,
+            system=system,
+            input_messages=input_messages,
+            input_tokens=find_replace_response.usage.input_tokens,
+            output_tokens=find_replace_response.usage.output_tokens,
             start_time=start_time,
             end_time=end_time,
         )
@@ -174,69 +170,64 @@ class Agent:
     async def __call__[R: Result](
         self,
         *,
-        messages: list[UserMessage | AssistantMessage],
-        model: AIModel,
+        messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
+        model: LLMModel,
         task_name: str,
         mode: Literal["generate"],
         response_model: type[R],
-        system: SystemMessage | None = None,
     ) -> R: ...
 
     @overload
     async def __call__[R: Result](
         self,
         *,
-        messages: list[UserMessage | AssistantMessage],
-        model: AIModel,
+        messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
+        model: LLMModel,
         task_name: str,
         mode: Literal["ask"],
         response_model: type[R] = TextResult,
-        system: SystemMessage | None = None,
     ) -> str: ...
 
     @overload
     async def __call__[R: Result](
         self,
         *,
-        messages: list[UserMessage | AssistantMessage],
-        model: AIModel,
+        messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
+        model: LLMModel,
         task_name: str,
         response_model: type[R],
         mode: Literal["refine"],
-        system: SystemMessage | None = None,
     ) -> R: ...
 
     @overload
     async def __call__[R: Result](
         self,
         *,
-        messages: list[UserMessage | AssistantMessage],
-        model: AIModel,
+        messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
+        model: LLMModel,
         task_name: str,
         response_model: type[R],
-        system: SystemMessage | None = None,
     ) -> R: ...
 
     async def __call__[R: Result](
         self,
         *,
-        messages: list[UserMessage | AssistantMessage],
-        model: AIModel,
+        messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage],
+        model: LLMModel,
         task_name: str,
         response_model: type[R] = TextResult,
         mode: Literal["generate", "ask", "refine"] = "generate",
-        system: SystemMessage | None = None,
     ) -> R | str:
         match mode:
             case "generate":
                 return await self.generate(
-                    messages=messages, model=model, task_name=task_name, response_model=response_model, system=system
+                    messages=messages, model=model, task_name=task_name, response_model=response_model
                 )
             case "ask":
-                return await self.ask(messages=messages, model=model, task_name=task_name, system=system)
+                return await self.ask(messages=messages, model=model, task_name=task_name)
             case "refine":
                 return await self.refine(
-                    messages=messages, model=model, task_name=task_name, response_model=response_model, system=system
+                    messages=messages, model=model, task_name=task_name, response_model=response_model
                 )
 
     def __apply_find_replace[R: Result](
@@ -249,12 +240,12 @@ class Agent:
         return response_model.model_validate_json(dumped_result)
 
     def __find_last_result[R: Result](
-        self, *, messages: list[UserMessage | AssistantMessage], response_model: type[R]
+        self, *, messages: list[LLMUserMessage | LLMAssistantMessage | LLMSystemMessage], response_model: type[R]
     ) -> R:
         for message in reversed(messages):
-            if isinstance(message, AssistantMessage):
+            if isinstance(message, LLMAssistantMessage):
                 try:
-                    return response_model.model_validate_json(message.text)
+                    return response_model.model_validate_json(message.parts[0].content)
                 except ValidationError:
                     continue
 

goose/_internal/conversation.py

@@ -1,38 +1,33 @@
 from typing import Self
 
+from aikernel import LLMAssistantMessage, LLMSystemMessage, LLMUserMessage
 from pydantic import BaseModel
 
-from goose._internal.result import Result
-from goose._internal.types.agent import AssistantMessage, SystemMessage, UserMessage
 from goose.errors import Honk
 
 
-class Conversation[R: Result](BaseModel):
-    user_messages: list[UserMessage]
-    assistant_messages: list[R | str]
-    context: SystemMessage | None = None
+class Conversation(BaseModel):
+    user_messages: list[LLMUserMessage]
+    assistant_messages: list[LLMAssistantMessage]
+    context: LLMSystemMessage | None = None
 
     @property
     def awaiting_response(self) -> bool:
         return len(self.user_messages) == len(self.assistant_messages)
 
-    def get_messages(self) -> list[UserMessage | AssistantMessage]:
-        messages: list[UserMessage | AssistantMessage] = []
+    def render(self) -> list[LLMSystemMessage | LLMUserMessage | LLMAssistantMessage]:
+        messages: list[LLMSystemMessage | LLMUserMessage | LLMAssistantMessage] = []
+        if self.context is not None:
+            messages.append(self.context)
+
         for message_index in range(len(self.user_messages)):
             message = self.assistant_messages[message_index]
-            if isinstance(message, str):
-                messages.append(AssistantMessage(text=message))
-            else:
-                messages.append(AssistantMessage(text=message.model_dump_json()))
-
+            messages.append(message)
             messages.append(self.user_messages[message_index])
 
         if len(self.assistant_messages) > len(self.user_messages):
             message = self.assistant_messages[-1]
-            if isinstance(message, str):
-                messages.append(AssistantMessage(text=message))
-            else:
-                messages.append(AssistantMessage(text=message.model_dump_json()))
+            messages.append(message)
 
         return messages
 

goose/_internal/flow.py

@@ -3,12 +3,12 @@ from contextlib import asynccontextmanager
 from types import CodeType
 from typing import Protocol, overload
 
-from goose._internal.agent import Agent, IAgentLogger
-from goose._internal.conversation import Conversation
-from goose._internal.result import Result
-from goose._internal.state import FlowArguments, FlowRun, get_current_flow_run, set_current_flow_run
-from goose._internal.store import IFlowRunStore, InMemoryFlowRunStore
-from goose.errors import Honk
+from ..errors import Honk
+from .agent import Agent, IAgentLogger
+from .conversation import Conversation
+from .result import Result
+from .state import FlowArguments, FlowRun, get_current_flow_run, set_current_flow_run
+from .store import IFlowRunStore, InMemoryFlowRunStore
 
 
 class IGenerator[FlowArgumentsT: FlowArguments](Protocol):
@@ -20,7 +20,7 @@ class IGenerator[FlowArgumentsT: FlowArguments](Protocol):
 class IAdapter[ResultT: Result](Protocol):
     __code__: CodeType
 
-    async def __call__(self, *, conversation: Conversation[ResultT], agent: Agent) -> ResultT: ...
+    async def __call__(self, *, conversation: Conversation, agent: Agent) -> ResultT: ...
 
 
 class Flow[FlowArgumentsT: FlowArguments]:

goose/_internal/state.py

@@ -2,12 +2,12 @@ import json
 from contextvars import ContextVar
 from typing import TYPE_CHECKING, Any, NewType, Self
 
+from aikernel import LLMAssistantMessage, LLMSystemMessage, LLMUserMessage
 from pydantic import BaseModel, ConfigDict
 
 from goose._internal.agent import Agent, IAgentLogger
 from goose._internal.conversation import Conversation
 from goose._internal.result import Result
-from goose._internal.types.agent import SystemMessage, UserMessage
 from goose.errors import Honk
 
 if TYPE_CHECKING:
@@ -20,55 +20,55 @@ class FlowArguments(BaseModel):
     model_config = ConfigDict(frozen=True)
 
 
-class NodeState[ResultT: Result](BaseModel):
+class NodeState(BaseModel):
     task_name: str
     index: int
-    conversation: Conversation[ResultT]
+    conversation: Conversation
     last_hash: int
 
     @property
-    def result(self) -> ResultT:
+    def raw_result(self) -> str:
         for message in reversed(self.conversation.assistant_messages):
-            if isinstance(message, Result):
-                return message
+            if self.__message_is_result(message):
+                return message.parts[0].content
 
         raise Honk("Node awaiting response, has no result")
 
-    def set_context(self, *, context: SystemMessage) -> Self:
+    def set_context(self, *, context: LLMSystemMessage) -> Self:
         self.conversation.context = context
         return self
 
     def add_result(
         self,
         *,
-        result: ResultT,
+        result: str,
         new_hash: int | None = None,
         overwrite: bool = False,
     ) -> Self:
         if overwrite and len(self.conversation.assistant_messages) > 0:
-            self.conversation.assistant_messages[-1] = result
+            self.conversation.assistant_messages[-1] = LLMAssistantMessage.from_text(result)
         else:
-            self.conversation.assistant_messages.append(result)
+            self.conversation.assistant_messages.append(LLMAssistantMessage.from_text(result))
         if new_hash is not None:
             self.last_hash = new_hash
         return self
 
     def add_answer(self, *, answer: str) -> Self:
-        self.conversation.assistant_messages.append(answer)
+        self.conversation.assistant_messages.append(LLMAssistantMessage.from_text(answer))
         return self
 
-    def add_user_message(self, *, message: UserMessage) -> Self:
+    def add_user_message(self, *, message: LLMUserMessage) -> Self:
         self.conversation.user_messages.append(message)
         return self
 
-    def edit_last_result(self, *, result: ResultT) -> Self:
+    def edit_last_result(self, *, result: str) -> Self:
         if len(self.conversation.assistant_messages) == 0:
             raise Honk("Node awaiting response, has no result")
 
         for message_index, message in enumerate(reversed(self.conversation.assistant_messages)):
-            if isinstance(message, Result):
+            if self.__message_is_result(message):
                 index = len(self.conversation.assistant_messages) - message_index - 1
-                self.conversation.assistant_messages[index] = result
+                self.conversation.assistant_messages[index] = LLMAssistantMessage.from_text(result)
                 return self
 
         raise Honk("Node awaiting response, has no result")
@@ -77,6 +77,13 @@ class NodeState[ResultT: Result](BaseModel):
         self.conversation.undo()
         return self
 
+    def __message_is_result(self, message: LLMAssistantMessage, /) -> bool:
+        try:
+            _ = json.loads(message.parts[0].content)
+            return True
+        except json.JSONDecodeError:
+            return False
+
 
 class FlowRun[FlowArgumentsT: FlowArguments]:
     def __init__(self, *, flow_arguments_model: type[FlowArgumentsT]) -> None:
@@ -109,38 +116,47 @@ class FlowRun[FlowArgumentsT: FlowArguments]:
 
         return self._flow_arguments
 
-    def get_all[R: Result](self, *, task: "Task[Any, R]") -> list[NodeState[R]]:
-        matching_nodes: list[NodeState[R]] = []
-        for key, node_state in self._node_states.items():
-            if key[0] == task.name:
-                matching_nodes.append(NodeState[task.result_type].model_validate_json(node_state))
-        return sorted(matching_nodes, key=lambda node: node.index)
-
-    def get[R: Result](self, *, task: "Task[Any, R]", index: int = 0) -> NodeState[R]:
+    def get_state(self, *, task: "Task[Any, Any]", index: int = 0) -> NodeState:
         if (existing_node_state := self._node_states.get((task.name, index))) is not None:
-            return NodeState[task.result_type].model_validate_json(existing_node_state)
+            return NodeState.model_validate_json(existing_node_state)
         else:
-            return NodeState[task.result_type](
+            return NodeState(
                 task_name=task.name,
                 index=index,
-                conversation=Conversation[task.result_type](user_messages=[], assistant_messages=[]),
+                conversation=Conversation(user_messages=[], assistant_messages=[]),
                 last_hash=0,
             )
 
-    def set_flow_arguments(self, flow_arguments: FlowArgumentsT, /) -> None:
-        self._flow_arguments = flow_arguments
-
-    def upsert_node_state(self, node_state: NodeState[Any], /) -> None:
-        key = (node_state.task_name, node_state.index)
-        self._node_states[key] = node_state.model_dump_json()
-
-    def get_next[R: Result](self, *, task: "Task[Any, R]") -> NodeState[R]:
+    def get_next_state(self, *, task: "Task[Any, Any]", index: int = 0) -> NodeState:
         if task.name not in self._last_requested_indices:
             self._last_requested_indices[task.name] = 0
         else:
             self._last_requested_indices[task.name] += 1
 
-        return self.get(task=task, index=self._last_requested_indices[task.name])
+        return self.get_state(task=task, index=self._last_requested_indices[task.name])
+
+    def get_all_results[R: Result](self, *, task: "Task[Any, R]") -> list[R]:
+        matching_nodes: list[NodeState] = []
+        for key, node_state in self._node_states.items():
+            if key[0] == task.name:
+                matching_nodes.append(NodeState.model_validate_json(node_state))
+
+        sorted_nodes = sorted(matching_nodes, key=lambda node: node.index)
+        return [task.result_type.model_validate_json(node.raw_result) for node in sorted_nodes]
+
+    def get_result[R: Result](self, *, task: "Task[Any, R]", index: int = 0) -> R:
+        if (existing_node_state := self._node_states.get((task.name, index))) is not None:
+            parsed_node_state = NodeState.model_validate_json(existing_node_state)
+            return task.result_type.model_validate_json(parsed_node_state.raw_result)
+        else:
+            raise Honk(f"No result found for task {task.name} at index {index}")
+
+    def set_flow_arguments(self, flow_arguments: FlowArgumentsT, /) -> None:
+        self._flow_arguments = flow_arguments
+
+    def upsert_node_state(self, node_state: NodeState, /) -> None:
+        key = (node_state.task_name, node_state.index)
+        self._node_states[key] = node_state.model_dump_json()
 
     def start(
         self,

goose/_internal/store.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import Protocol
 
-from goose._internal.state import SerializedFlowRun
+from .state import SerializedFlowRun
 
 
 class IFlowRunStore(Protocol):

goose/_internal/task.py

@@ -2,12 +2,12 @@ import hashlib
 from collections.abc import Awaitable, Callable
 from typing import Any, overload
 
+from aikernel import LLMModel, LLMSystemMessage, LLMUserMessage
 from pydantic import BaseModel
 
-from goose._internal.agent import Agent, AIModel
+from goose._internal.agent import Agent
 from goose._internal.result import Result
 from goose._internal.state import FlowRun, NodeState, get_current_flow_run
-from goose._internal.types.agent import SystemMessage, UserMessage
 from goose.errors import Honk
 
 
@@ -18,7 +18,7 @@ class Task[**P, R: Result]:
         /,
         *,
         retries: int = 0,
-        refinement_model: AIModel = AIModel.GEMINI_FLASH,
+        refinement_model: LLMModel = LLMModel.GEMINI_2_0_FLASH,
     ) -> None:
         self._generator = generator
         self._retries = retries
@@ -35,28 +35,32 @@ class Task[**P, R: Result]:
     def name(self) -> str:
         return self._generator.__name__
 
-    async def generate(self, state: NodeState[R], *args: P.args, **kwargs: P.kwargs) -> R:
+    async def generate(self, state: NodeState, *args: P.args, **kwargs: P.kwargs) -> R:
         state_hash = self.__hash_task_call(*args, **kwargs)
         if state_hash != state.last_hash:
             result = await self._generator(*args, **kwargs)
-            state.add_result(result=result, new_hash=state_hash, overwrite=True)
+            state.add_result(result=result.model_dump_json(), new_hash=state_hash, overwrite=True)
             return result
         else:
-            return state.result
+            return self.result_type.model_validate_json(state.raw_result)
 
-    async def ask(self, *, user_message: UserMessage, context: SystemMessage | None = None, index: int = 0) -> str:
+    async def ask(
+        self, *, user_message: LLMUserMessage, context: LLMSystemMessage | None = None, index: int = 0
+    ) -> str:
         flow_run = self.__get_current_flow_run()
-        node_state = flow_run.get(task=self, index=index)
+        node_state = flow_run.get_state(task=self, index=index)
 
         if len(node_state.conversation.assistant_messages) == 0:
             raise Honk("Cannot ask about a task that has not been initially generated")
 
+        if context is not None:
+            node_state.set_context(context=context)
         node_state.add_user_message(message=user_message)
+
         answer = await flow_run.agent(
-            messages=node_state.conversation.get_messages(),
+            messages=node_state.conversation.render(),
             model=self._refinement_model,
             task_name=f"ask--{self.name}",
-            system=context,
             mode="ask",
         )
         node_state.add_answer(answer=answer)
@@ -67,12 +71,12 @@ class Task[**P, R: Result]:
     async def refine(
         self,
         *,
-        user_message: UserMessage,
-        context: SystemMessage | None = None,
+        user_message: LLMUserMessage,
+        context: LLMSystemMessage | None = None,
         index: int = 0,
     ) -> R:
         flow_run = self.__get_current_flow_run()
-        node_state = flow_run.get(task=self, index=index)
+        node_state = flow_run.get_state(task=self, index=index)
 
         if len(node_state.conversation.assistant_messages) == 0:
             raise Honk("Cannot refine a task that has not been initially generated")
@@ -82,33 +86,32 @@ class Task[**P, R: Result]:
         node_state.add_user_message(message=user_message)
 
         result = await flow_run.agent(
-            messages=node_state.conversation.get_messages(),
+            messages=node_state.conversation.render(),
             model=self._refinement_model,
             task_name=f"refine--{self.name}",
-            system=context,
             response_model=self.result_type,
             mode="refine",
         )
-        node_state.add_result(result=result)
+        node_state.add_result(result=result.model_dump_json())
         flow_run.upsert_node_state(node_state)
 
         return result
 
     def edit(self, *, result: R, index: int = 0) -> None:
         flow_run = self.__get_current_flow_run()
-        node_state = flow_run.get(task=self, index=index)
-        node_state.edit_last_result(result=result)
+        node_state = flow_run.get_state(task=self, index=index)
+        node_state.edit_last_result(result=result.model_dump_json())
         flow_run.upsert_node_state(node_state)
 
     def undo(self, *, index: int = 0) -> None:
         flow_run = self.__get_current_flow_run()
-        node_state = flow_run.get(task=self, index=index)
+        node_state = flow_run.get_state(task=self, index=index)
         node_state.undo()
         flow_run.upsert_node_state(node_state)
 
     async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
         flow_run = self.__get_current_flow_run()
-        node_state = flow_run.get_next(task=self)
+        node_state = flow_run.get_next_state(task=self)
         result = await self.generate(node_state, *args, **kwargs)
         flow_run.upsert_node_state(node_state)
         return result
@@ -151,14 +154,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: AIModel = AIModel.GEMINI_FLASH
+    *, retries: int = 0, refinement_model: LLMModel = LLMModel.GEMINI_2_0_FLASH
 ) -> 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: AIModel = AIModel.GEMINI_FLASH,
+    refinement_model: LLMModel = LLMModel.GEMINI_2_0_FLASH,
 ) -> Task[P, R] | Callable[[Callable[P, Awaitable[R]]], Task[P, R]]:
     if generator is None:
 

goose/_internal/types/telemetry.py

@@ -2,10 +2,9 @@ import json
 from datetime import datetime
 from typing import ClassVar, TypedDict
 
+from aikernel import LiteLLMMessage, LLMModel
 from pydantic import BaseModel, computed_field
 
-from ..types.agent import AIModel, LLMMessage
-
 
 class AgentResponseDump(TypedDict):
     run_id: str
@@ -26,30 +25,30 @@ class AgentResponseDump(TypedDict):
 
 
 class AgentResponse[R: BaseModel | str](BaseModel):
-    INPUT_DOLLARS_PER_MILLION_TOKENS: ClassVar[dict[AIModel, float]] = {
-        AIModel.VERTEX_FLASH_8B: 0.30,
-        AIModel.VERTEX_FLASH: 0.15,
-        AIModel.VERTEX_PRO: 5.00,
-        AIModel.GEMINI_FLASH_8B: 0.30,
-        AIModel.GEMINI_FLASH: 0.15,
-        AIModel.GEMINI_PRO: 5.00,
+    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,
     }
-    OUTPUT_DOLLARS_PER_MILLION_TOKENS: ClassVar[dict[AIModel, float]] = {
-        AIModel.VERTEX_FLASH_8B: 0.30,
-        AIModel.VERTEX_FLASH: 0.15,
-        AIModel.VERTEX_PRO: 5.00,
-        AIModel.GEMINI_FLASH_8B: 0.30,
-        AIModel.GEMINI_FLASH: 0.15,
-        AIModel.GEMINI_PRO: 5.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,
     }
 
     response: R
     run_id: str
     flow_name: str
     task_name: str
-    model: AIModel
-    system: LLMMessage | None = None
-    input_messages: list[LLMMessage]
+    model: LLMModel
+    system: LiteLLMMessage | None = None
+    input_messages: list[LiteLLMMessage]
     input_tokens: int
     output_tokens: int
     start_time: datetime

{goose_py-0.10.2.dist-info → goose_py-0.11.1.dist-info}/METADATA

@@ -1,11 +1,11 @@
 Metadata-Version: 2.4
 Name: goose-py
-Version: 0.10.2
+Version: 0.11.1
 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.8
 Requires-Dist: jsonpath-ng>=1.7.0
-Requires-Dist: litellm>=1.56.5
 Requires-Dist: pydantic>=2.8.2
 Description-Content-Type: text/markdown
 

goose_py-0.11.1.dist-info/RECORD

@@ -0,0 +1,18 @@
+goose/__init__.py,sha256=Muw7HCImZHk3kLCTWhV9Lg-Sfmhnwf_Tae-zCj7woyY,338
+goose/errors.py,sha256=-0OyZQJWYTRw5YgnCB2_uorVaUsL6Z0QYQO2FqzCiyg,32
+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=rAE7P0DFl-r-omjrxR2EzkJUKQWbaPM0qM3sxTr70ew,8366
+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=jGKKZKeDcKaCY-uqgXLOnMNwyzR5qEH5m9afOtixsLk,7923
+goose/_internal/store.py,sha256=tWmKfa1-yq1jU6lT3l6kSOmVt2m3H7I1xLMTrxnUDI8,889
+goose/_internal/task.py,sha256=X_eRZxZlf6SwyvF1nIyjoneyqD_TISXqESyxluk63mE,6416
+goose/_internal/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+goose/_internal/types/telemetry.py,sha256=wDQz1C60KxlftC0aQCXIjswrSYGI1KNaM2wtnlP1Q7k,3823
+goose_py-0.11.1.dist-info/METADATA,sha256=mAoNCBZ4vauZDCSqhCRS4zVSZMb4sXGS08epirIUiMY,442
+goose_py-0.11.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+goose_py-0.11.1.dist-info/RECORD,,

goose/_internal/types/agent.py

@@ -1,101 +0,0 @@
-import base64
-from enum import StrEnum
-from typing import Literal, NotRequired, TypedDict
-
-from pydantic import BaseModel
-
-
-class AIModel(StrEnum):
-    # vertex (production Google, requires GCP environment)
-    VERTEX_PRO = "vertex_ai/gemini-1.5-pro"
-    VERTEX_FLASH = "vertex_ai/gemini-1.5-flash"
-    VERTEX_FLASH_8B = "vertex_ai/gemini-1.5-flash-8b"
-    VERTEX_FLASH_2_0 = "vertex_ai/gemini-2.0-flash"
-
-    # gemini (publicly available, no GCP environment required)
-    GEMINI_PRO = "gemini/gemini-1.5-pro"
-    GEMINI_FLASH = "gemini/gemini-1.5-flash"
-    GEMINI_FLASH_8B = "gemini/gemini-1.5-flash-8b"
-    GEMINI_FLASH_2_0 = "gemini/gemini-2.0-flash"
-
-
-class ContentType(StrEnum):
-    # text
-    TEXT = "text/plain"
-
-    # images
-    JPEG = "image/jpeg"
-    PNG = "image/png"
-    WEBP = "image/webp"
-
-    # audio
-    MP3 = "audio/mp3"
-    WAV = "audio/wav"
-
-    # files
-    PDF = "application/pdf"
-
-
-class LLMTextMessagePart(TypedDict):
-    type: Literal["text"]
-    text: str
-
-
-class LLMMediaMessagePart(TypedDict):
-    type: Literal["image_url"]
-    image_url: str
-
-
-class CacheControl(TypedDict):
-    type: Literal["ephemeral"]
-
-
-class LLMMessage(TypedDict):
-    role: Literal["user", "assistant", "system"]
-    content: list[LLMTextMessagePart | LLMMediaMessagePart]
-    cache_control: NotRequired[CacheControl]
-
-
-class MessagePart(BaseModel):
-    content: str
-    content_type: ContentType = ContentType.TEXT
-
-    @classmethod
-    def from_media(cls, *, content: bytes, content_type: ContentType) -> "MessagePart":
-        return cls(content=base64.b64encode(content).decode(), content_type=content_type)
-
-    def render(self) -> LLMTextMessagePart | LLMMediaMessagePart:
-        if self.content_type == ContentType.TEXT:
-            return {"type": "text", "text": self.content}
-        else:
-            return {"type": "image_url", "image_url": f"data:{self.content_type};base64,{self.content}"}
-
-
-class UserMessage(BaseModel):
-    parts: list[MessagePart]
-
-    def render(self) -> LLMMessage:
-        content: LLMMessage = {
-            "role": "user",
-            "content": [part.render() for part in self.parts],
-        }
-        if any(part.content_type != ContentType.TEXT for part in self.parts):
-            content["cache_control"] = {"type": "ephemeral"}
-        return content
-
-
-class AssistantMessage(BaseModel):
-    text: str
-
-    def render(self) -> LLMMessage:
-        return {"role": "assistant", "content": [{"type": "text", "text": self.text}]}
-
-
-class SystemMessage(BaseModel):
-    parts: list[MessagePart]
-
-    def render(self) -> LLMMessage:
-        return {
-            "role": "system",
-            "content": [part.render() for part in self.parts],
-        }

goose/agent.py

@@ -1,26 +0,0 @@
-from ._internal.agent import AgentResponse, IAgentLogger
-from ._internal.types.agent import (
-    AIModel,
-    AssistantMessage,
-    ContentType,
-    LLMMediaMessagePart,
-    LLMMessage,
-    LLMTextMessagePart,
-    MessagePart,
-    SystemMessage,
-    UserMessage,
-)
-
-__all__ = [
-    "AgentResponse",
-    "AIModel",
-    "IAgentLogger",
-    "AssistantMessage",
-    "LLMMediaMessagePart",
-    "LLMMessage",
-    "LLMTextMessagePart",
-    "SystemMessage",
-    "MessagePart",
-    "ContentType",
-    "UserMessage",
-]

goose_py-0.10.2.dist-info/RECORD

@@ -1,20 +0,0 @@
-goose/__init__.py,sha256=wjGDgWzKcD6S8loVr0n-rLCpRwg-ZKAixcUaw1wobMc,243
-goose/agent.py,sha256=u6daAnn4fPgP4Jk9cHANyCEku3RmUqKLdqtyGSr8ljI,510
-goose/errors.py,sha256=-0OyZQJWYTRw5YgnCB2_uorVaUsL6Z0QYQO2FqzCiyg,32
-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=yNkX0lZueKiGxYsSwbZRNn88HfDBYbVpJLDfyr6Var0,8893
-goose/_internal/conversation.py,sha256=I0Ru5D7piOOQlmFhUpwEeQUQxOVq59yyHEDrmYhwPMI,1695
-goose/_internal/flow.py,sha256=05U2f5i8ofQWDjghhomwuuEPMk-ftzXn7BVl_s7pIf8,4203
-goose/_internal/result.py,sha256=vtJMfBxb9skfl8st2tn4hBmEq6qmXiJTme_B5QTgu2M,538
-goose/_internal/state.py,sha256=U4gM0K4MAlRFTpqenCYHX9TYGuhWVKIfa4yBeZ9Qc9s,7090
-goose/_internal/store.py,sha256=GMW0wBpxESmRBLfL_lFKEi9x2P6Wd6-gZ7AWjWBTUmA,904
-goose/_internal/task.py,sha256=MXBVLepFSphkSbGgzh7U7QMRXDvz_4MNet_EsrlAKTQ,6244
-goose/_internal/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-goose/_internal/types/agent.py,sha256=g0KD-aPWZlUGBx72AwQd3LeniFxHATeflZ7191QjFZA,2696
-goose/_internal/types/telemetry.py,sha256=7zeqyDDxf95puirNM6Gr9VFuxoDshXcV1__V0tiMswE,3663
-goose_py-0.10.2.dist-info/METADATA,sha256=VoLeEDnKj2PfjzVyJCt8eSQQygbFyWFihpADvJ4DTLw,442
-goose_py-0.10.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-goose_py-0.10.2.dist-info/RECORD,,