goose-py 0.10.2__tar.gz → 0.11.1__tar.gz

{goose_py-0.10.2 → goose_py-0.11.1}/PKG-INFO

@@ -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/goose/__init__.py

@@ -0,0 +1,7 @@
+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", "AgentResponse"]

{goose_py-0.10.2 → goose_py-0.11.1}/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_py-0.10.2 → goose_py-0.11.1}/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_py-0.10.2 → goose_py-0.11.1}/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_py-0.10.2 → goose_py-0.11.1}/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_py-0.10.2 → goose_py-0.11.1}/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):