goose-py 0.9.16__tar.gz → 0.10.0__tar.gz

{goose_py-0.9.16 → goose_py-0.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: goose-py
-Version: 0.9.16
+Version: 0.10.0
 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

goose_py-0.10.0/goose/_internal/agent.py

@@ -0,0 +1,245 @@
+import logging
+from datetime import datetime
+from typing import Any, Literal, Protocol, overload
+
+from litellm import acompletion
+from pydantic import ValidationError
+
+from goose._internal.types.telemetry import AgentResponse
+from goose.errors import Honk
+
+from .result import FindReplaceResponse, Result, TextResult
+from .types.agent import AIModel, LLMMessage
+
+
+class IAgentLogger(Protocol):
+    async def __call__(self, *, response: AgentResponse[Any]) -> None: ...
+
+
+class Agent:
+    def __init__(
+        self,
+        *,
+        flow_name: str,
+        run_id: str,
+        logger: IAgentLogger | None = None,
+    ) -> None:
+        self.flow_name = flow_name
+        self.run_id = run_id
+        self.logger = logger
+
+    async def generate[R: Result](
+        self,
+        *,
+        messages: list[LLMMessage],
+        model: AIModel,
+        task_name: str,
+        response_model: type[R] = TextResult,
+        system: LLMMessage | None = None,
+    ) -> R:
+        start_time = datetime.now()
+        if system is not None:
+            messages.insert(0, system)
+
+        if response_model is TextResult:
+            response = await acompletion(model=model.value, messages=messages)
+            parsed_response = response_model.model_validate({"text": response.choices[0].message.content})
+        else:
+            response = await acompletion(
+                model=model.value,
+                messages=messages,
+                response_format=response_model,
+            )
+            parsed_response = response_model.model_validate_json(response.choices[0].message.content)
+
+        end_time = datetime.now()
+        agent_response = AgentResponse(
+            response=parsed_response,
+            run_id=self.run_id,
+            flow_name=self.flow_name,
+            task_name=task_name,
+            model=model,
+            system=system,
+            input_messages=messages,
+            input_tokens=response.usage.prompt_tokens,
+            output_tokens=response.usage.completion_tokens,
+            start_time=start_time,
+            end_time=end_time,
+        )
+
+        if self.logger is not None:
+            await self.logger(response=agent_response)
+        else:
+            logging.info(agent_response.model_dump())
+
+        return parsed_response
+
+    async def ask(
+        self, *, messages: list[LLMMessage], model: AIModel, task_name: str, system: LLMMessage | None = None
+    ) -> str:
+        start_time = datetime.now()
+
+        if system is not None:
+            messages.insert(0, system)
+        response = await acompletion(model=model.value, messages=messages)
+
+        end_time = datetime.now()
+        agent_response = AgentResponse(
+            response=response.choices[0].message.content,
+            run_id=self.run_id,
+            flow_name=self.flow_name,
+            task_name=task_name,
+            model=model,
+            system=system,
+            input_messages=messages,
+            input_tokens=response.usage.prompt_tokens,
+            output_tokens=response.usage.completion_tokens,
+            start_time=start_time,
+            end_time=end_time,
+        )
+
+        if self.logger is not None:
+            await self.logger(response=agent_response)
+        else:
+            logging.info(agent_response.model_dump())
+
+        return response.choices[0].message.content
+
+    async def refine[R: Result](
+        self,
+        *,
+        messages: list[LLMMessage],
+        model: AIModel,
+        task_name: str,
+        response_model: type[R],
+        system: LLMMessage | None = None,
+    ) -> R:
+        start_time = datetime.now()
+
+        if system is not None:
+            messages.insert(0, system)
+
+        find_replace_response = await acompletion(
+            model=model.value, messages=messages, response_format=FindReplaceResponse
+        )
+        parsed_find_replace_response = FindReplaceResponse.model_validate_json(
+            find_replace_response.choices[0].message.content
+        )
+
+        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=system,
+            input_messages=messages,
+            input_tokens=find_replace_response.usage.prompt_tokens,
+            output_tokens=find_replace_response.usage.completion_tokens,
+            start_time=start_time,
+            end_time=end_time,
+        )
+
+        if self.logger is not None:
+            await self.logger(response=agent_response)
+        else:
+            logging.info(agent_response.model_dump())
+
+        refined_response = self.__apply_find_replace(
+            result=self.__find_last_result(messages=messages, response_model=response_model),
+            find_replace_response=parsed_find_replace_response,
+            response_model=response_model,
+        )
+
+        return refined_response
+
+    @overload
+    async def __call__[R: Result](
+        self,
+        *,
+        messages: list[LLMMessage],
+        model: AIModel,
+        task_name: str,
+        mode: Literal["generate"],
+        response_model: type[R],
+        system: LLMMessage | None = None,
+    ) -> R: ...
+
+    @overload
+    async def __call__[R: Result](
+        self,
+        *,
+        messages: list[LLMMessage],
+        model: AIModel,
+        task_name: str,
+        mode: Literal["ask"],
+        response_model: type[R] = TextResult,
+        system: LLMMessage | None = None,
+    ) -> str: ...
+
+    @overload
+    async def __call__[R: Result](
+        self,
+        *,
+        messages: list[LLMMessage],
+        model: AIModel,
+        task_name: str,
+        response_model: type[R],
+        mode: Literal["refine"],
+        system: LLMMessage | None = None,
+    ) -> R: ...
+
+    @overload
+    async def __call__[R: Result](
+        self,
+        *,
+        messages: list[LLMMessage],
+        model: AIModel,
+        task_name: str,
+        response_model: type[R],
+        system: LLMMessage | None = None,
+    ) -> R: ...
+
+    async def __call__[R: Result](
+        self,
+        *,
+        messages: list[LLMMessage],
+        model: AIModel,
+        task_name: str,
+        response_model: type[R] = TextResult,
+        mode: Literal["generate", "ask", "refine"] = "generate",
+        system: LLMMessage | 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
+                )
+            case "ask":
+                return await self.ask(messages=messages, model=model, task_name=task_name, system=system)
+            case "refine":
+                return await self.refine(
+                    messages=messages, model=model, task_name=task_name, response_model=response_model, system=system
+                )
+
+    def __apply_find_replace[R: Result](
+        self, *, result: R, find_replace_response: FindReplaceResponse, response_model: type[R]
+    ) -> R:
+        dumped_result = result.model_dump_json()
+        for replacement in find_replace_response.replacements:
+            dumped_result = dumped_result.replace(replacement.find, replacement.replace)
+
+        return response_model.model_validate_json(dumped_result)
+
+    def __find_last_result[R: Result](self, *, messages: list[LLMMessage], response_model: type[R]) -> R:
+        for message in reversed(messages):
+            if message["role"] == "assistant":
+                try:
+                    only_part = message["content"][0]
+                    if only_part["type"] == "text":
+                        return response_model.model_validate_json(only_part["text"])
+                except ValidationError:
+                    continue
+
+        raise Honk("No last result found, failed to refine")

goose_py-0.10.0/goose/_internal/result.py

@@ -0,0 +1,20 @@
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class Result(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+
+class TextResult(Result):
+    text: str
+
+
+class Replacement(BaseModel):
+    find: str = Field(description="Text to find, to be replaced with `replace`")
+    replace: str = Field(description="Text to replace `find` with")
+
+
+class FindReplaceResponse(BaseModel):
+    replacements: list[Replacement] = Field(
+        description="List of replacements to make in the previous result to satisfy the user's request"
+    )

{goose_py-0.9.16 → goose_py-0.10.0}/goose/_internal/task.py

@@ -57,11 +57,12 @@ class Task[**P, R: Result]:
             model=self._refinement_model,
             task_name=f"ask--{self.name}",
             system=context.render() if context is not None else None,
+            mode="ask",
         )
-        node_state.add_answer(answer=answer.text)
+        node_state.add_answer(answer=answer)
         flow_run.upsert_node_state(node_state)
 
-        return answer.text
+        return answer
 
     async def refine(
         self,

goose_py-0.10.0/goose/_internal/types/telemetry.py

@@ -0,0 +1,113 @@
+import json
+from datetime import datetime
+from typing import ClassVar, TypedDict
+
+from pydantic import BaseModel, computed_field
+
+from ..types.agent import AIModel, LLMMessage
+
+
+class AgentResponseDump(TypedDict):
+    run_id: str
+    flow_name: str
+    task_name: str
+    model: str
+    system_message: str
+    input_messages: list[str]
+    output_message: str
+    input_cost: float
+    output_cost: float
+    total_cost: float
+    input_tokens: int
+    output_tokens: int
+    start_time: datetime
+    end_time: datetime
+    duration_ms: int
+
+
+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,
+    }
+    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,
+    }
+
+    response: R
+    run_id: str
+    flow_name: str
+    task_name: str
+    model: AIModel
+    system: LLMMessage | None = None
+    input_messages: list[LLMMessage]
+    input_tokens: int
+    output_tokens: int
+    start_time: datetime
+    end_time: datetime
+
+    @computed_field
+    @property
+    def duration_ms(self) -> int:
+        return int((self.end_time - self.start_time).total_seconds() * 1000)
+
+    @computed_field
+    @property
+    def input_cost(self) -> float:
+        return self.INPUT_DOLLARS_PER_MILLION_TOKENS[self.model] * self.input_tokens / 1_000_000
+
+    @computed_field
+    @property
+    def output_cost(self) -> float:
+        return self.OUTPUT_DOLLARS_PER_MILLION_TOKENS[self.model] * self.output_tokens / 1_000_000
+
+    @computed_field
+    @property
+    def total_cost(self) -> float:
+        return self.input_cost + self.output_cost
+
+    def minimized_dump(self) -> AgentResponseDump:
+        if self.system is None:
+            minimized_system_message = ""
+        else:
+            minimized_system_message = self.system
+            for part in minimized_system_message["content"]:
+                if part["type"] == "image_url":
+                    part["image_url"] = "__MEDIA__"
+            minimized_system_message = json.dumps(minimized_system_message)
+
+        minimized_input_messages = [message for message in self.input_messages]
+        for message in minimized_input_messages:
+            for part in message["content"]:
+                if part["type"] == "image_url":
+                    part["image_url"] = "__MEDIA__"
+        minimized_input_messages = [json.dumps(message) for message in minimized_input_messages]
+
+        output_message = self.response.model_dump_json() if isinstance(self.response, BaseModel) else self.response
+
+        return {
+            "run_id": self.run_id,
+            "flow_name": self.flow_name,
+            "task_name": self.task_name,
+            "model": self.model.value,
+            "system_message": minimized_system_message,
+            "input_messages": minimized_input_messages,
+            "output_message": output_message,
+            "input_tokens": self.input_tokens,
+            "output_tokens": self.output_tokens,
+            "input_cost": self.input_cost,
+            "output_cost": self.output_cost,
+            "total_cost": self.total_cost,
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "duration_ms": self.duration_ms,
+        }

{goose_py-0.9.16 → goose_py-0.10.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "goose-py"
-version = "0.9.16"
+version = "0.10.0"
 description = "A tool for AI workflows based on human-computer collaboration and structured output."
 readme = "README.md"
 authors = [
@@ -24,6 +24,7 @@ packages = ["goose"]
 
 [dependency-groups]
 dev = [
+    "ipykernel>=6.29.5",
     "pyright>=1.1.393",
     "pytest>=8.3.4",
     "pytest-asyncio>=0.25.3",