goose-py 0.6.0__tar.gz → 0.7.1__tar.gz

goose_py-0.7.1/.envrc

@@ -0,0 +1 @@
+dotenv

goose_py-0.7.1/.github/workflows/publish.yml

@@ -0,0 +1,48 @@
+name: CI/CD
+
+on:
+  push:
+    branches:
+      - main
+    tags:
+      - "v*"
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+
+    permissions:
+      id-token: write
+    
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version-file: .python-version
+
+      - name: Setup UV
+        uses: astral-sh/setup-uv@v5
+        with:
+          version: "0.5.25"
+
+      - name: Initialize environment
+        run: uv sync --all-extras --dev
+
+      - name: Run tests
+        run: uv run pytest
+        env:
+          GEMINI_API_KEY: ${{ secrets.GEMINI_API_KEY }}
+
+      - name: Build package
+        run: uv build
+
+      - name: Publish package
+        run: uv publish
+        continue-on-error: true

goose_py-0.7.1/.gitignore

@@ -0,0 +1,5 @@
+__pycache__
+.env
+poetry.lock
+notebooks
+dist

goose_py-0.7.1/.python-version

@@ -0,0 +1 @@
+3.12

goose_py-0.7.1/.stubs/jsonpath_ng/__init__.pyi

@@ -0,0 +1,14 @@
+from typing import Any
+
+class Fields:
+    fields: tuple[str, ...]
+
+class DatumInContext:
+    value: Any
+    fields: Fields
+    context: "DatumInContext"
+
+class Jsonpath:
+    def find(self, /, value: Any) -> list[DatumInContext]: ...
+
+def parse(jsonpath: str, /) -> Jsonpath: ...

goose_py-0.7.1/.stubs/litellm/__init__.pyi

@@ -0,0 +1,61 @@
+from typing import Any, Literal, NotRequired, TypedDict
+
+_LiteLLMGeminiModel = Literal[
+    "vertex_ai/gemini-1.5-flash",
+    "vertex_ai/gemini-1.5-pro",
+    "vertex_ai/gemini-1.5-flash-8b",
+]
+_MessageRole = Literal["system", "user", "assistant"]
+
+class _LiteLLMTextMessageContent(TypedDict):
+    type: Literal["text"]
+    text: str
+
+class _LiteLLMMediaMessageContent(TypedDict):
+    type: Literal["image_url"]
+    image_url: str
+
+class _LiteLLMCacheControl(TypedDict):
+    type: Literal["ephemeral"]
+
+class _LiteLLMMessage(TypedDict):
+    role: _MessageRole
+    content: list[_LiteLLMTextMessageContent | _LiteLLMMediaMessageContent]
+    cache_control: NotRequired[_LiteLLMCacheControl]
+
+class _LiteLLMResponseFormat(TypedDict):
+    type: Literal["json_object"]
+    response_schema: dict[str, Any]  # must be a valid JSON schema
+    enforce_validation: NotRequired[bool]
+
+class _LiteLLMModelResponseChoiceMessage:
+    role: Literal["assistant"]
+    content: str
+
+class _LiteLLMModelResponseChoice:
+    finish_reason: Literal["stop"]
+    index: int
+    message: _LiteLLMModelResponseChoiceMessage
+
+class _LiteLLMUsage:
+    completion_tokens: int
+    prompt_tokens: int
+    total_tokens: int
+
+class ModelResponse:
+    id: str
+    created: int
+    model: _LiteLLMGeminiModel
+    object: Literal["chat.completion"]
+    system_fingerprint: str | None
+    choices: list[_LiteLLMModelResponseChoice]
+    usage: _LiteLLMUsage
+
+async def acompletion(
+    *,
+    model: _LiteLLMGeminiModel,
+    messages: list[_LiteLLMMessage],
+    response_format: _LiteLLMResponseFormat | None = None,
+    max_tokens: int | None = None,
+    temperature: float = 1.0,
+) -> ModelResponse: ...

goose_py-0.7.1/Makefile

@@ -0,0 +1,9 @@
+test:
+	uv run pytest tests
+lint:
+	uv run ruff check .
+	uv run ruff format .
+	uv run pyright .
+publish:
+	uv build
+	uv publish

goose_py-0.7.1/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: goose-py
+Version: 0.7.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: jsonpath-ng>=1.7.0
+Requires-Dist: litellm>=1.56.5
+Requires-Dist: pydantic>=2.8.2
+Description-Content-Type: text/markdown
+
+# Goose
+
+Docs to come.

goose_py-0.7.1/goose/__init__.py

@@ -0,0 +1,4 @@
+from goose._internal.agent import Agent
+from goose._internal.flow import flow
+from goose._internal.result import Result, TextResult
+from goose._internal.task import task

goose_py-0.6.0/goose/_agent.py → goose_py-0.7.1/goose/_internal/agent.py

@@ -6,8 +6,13 @@ from typing import Any, ClassVar, Protocol, TypedDict
 from litellm import acompletion
 from pydantic import BaseModel, computed_field
 
-from goose._result import Result, TextResult
-from goose.types.agent import AssistantMessage, GeminiModel, SystemMessage, UserMessage
+from goose._internal.result import Result, TextResult
+from goose._internal.types.agent import (
+    AssistantMessage,
+    GeminiModel,
+    SystemMessage,
+    UserMessage,
+)
 
 
 class AgentResponseDump(TypedDict):
@@ -60,20 +65,12 @@ class AgentResponse[R: BaseModel | str](BaseModel):
     @computed_field
     @property
     def input_cost(self) -> float:
-        return (
-            self.INPUT_CENTS_PER_MILLION_TOKENS[self.model]
-            * self.input_tokens
-            / 1_000_000
-        )
+        return self.INPUT_CENTS_PER_MILLION_TOKENS[self.model] * self.input_tokens / 1_000_000
 
     @computed_field
     @property
     def output_cost(self) -> float:
-        return (
-            self.OUTPUT_CENTS_PER_MILLION_TOKENS[self.model]
-            * self.output_tokens
-            / 1_000_000
-        )
+        return self.OUTPUT_CENTS_PER_MILLION_TOKENS[self.model] * self.output_tokens / 1_000_000
 
     @computed_field
     @property
@@ -95,15 +92,9 @@ class AgentResponse[R: BaseModel | str](BaseModel):
             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
-        )
+        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,
@@ -156,9 +147,7 @@ class Agent:
 
         if response_model is TextResult:
             response = await acompletion(model=model.value, messages=rendered_messages)
-            parsed_response = response_model.model_validate(
-                {"text": response.choices[0].message.content}
-            )
+            parsed_response = response_model.model_validate({"text": response.choices[0].message.content})
         else:
             response = await acompletion(
                 model=model.value,
@@ -169,9 +158,7 @@ class Agent:
                     "enforce_validation": True,
                 },
             )
-            parsed_response = response_model.model_validate_json(
-                response.choices[0].message.content
-            )
+            parsed_response = response_model.model_validate_json(response.choices[0].message.content)
 
         end_time = datetime.now()
         agent_response = AgentResponse(

goose_py-0.6.0/goose/_conversation.py → goose_py-0.7.1/goose/_internal/conversation.py

@@ -1,7 +1,12 @@
 from pydantic import BaseModel
 
-from goose._result import Result
-from goose.types.agent import AssistantMessage, LLMMessage, SystemMessage, UserMessage
+from goose._internal.result import Result
+from goose._internal.types.agent import (
+    AssistantMessage,
+    LLMMessage,
+    SystemMessage,
+    UserMessage,
+)
 
 
 class Conversation[R: Result](BaseModel):
@@ -19,18 +24,10 @@ class Conversation[R: Result](BaseModel):
             messages.append(self.context.render())
 
         for message_index in range(len(self.user_messages)):
-            messages.append(
-                AssistantMessage(
-                    text=self.result_messages[message_index].model_dump_json()
-                ).render()
-            )
+            messages.append(AssistantMessage(text=self.result_messages[message_index].model_dump_json()).render())
             messages.append(self.user_messages[message_index].render())
 
         if len(self.result_messages) > len(self.user_messages):
-            messages.append(
-                AssistantMessage(
-                    text=self.result_messages[-1].model_dump_json()
-                ).render()
-            )
+            messages.append(AssistantMessage(text=self.result_messages[-1].model_dump_json()).render())
 
         return messages

goose_py-0.6.0/goose/_flow.py → goose_py-0.7.1/goose/_internal/flow.py

@@ -1,21 +1,20 @@
+from collections.abc import AsyncIterator, Awaitable, Callable
 from contextlib import asynccontextmanager
 from types import CodeType
-from typing import AsyncIterator, Awaitable, Callable, Protocol, overload
+from typing import Protocol, overload
 
-from goose._agent import Agent, IAgentLogger
-from goose._conversation import Conversation
-from goose._result import Result
-from goose._state import FlowRun, get_current_flow_run, set_current_flow_run
-from goose._store import IFlowRunStore, InMemoryFlowRunStore
+from goose._internal.agent import Agent, IAgentLogger
+from goose._internal.conversation import Conversation
+from goose._internal.result import Result
+from goose._internal.state import FlowRun, get_current_flow_run, set_current_flow_run
+from goose._internal.store import IFlowRunStore, InMemoryFlowRunStore
 from goose.errors import Honk
 
 
 class IAdapter[ResultT: Result](Protocol):
     __code__: CodeType
 
-    async def __call__(
-        self, *, conversation: Conversation[ResultT], agent: Agent
-    ) -> ResultT: ...
+    async def __call__(self, *, conversation: Conversation[ResultT], agent: Agent) -> ResultT: ...
 
 
 class Flow[**P]:

goose_py-0.6.0/goose/_state.py → goose_py-0.7.1/goose/_internal/state.py

@@ -4,18 +4,18 @@ from typing import TYPE_CHECKING, Any, Self
 
 from pydantic import BaseModel
 
-from goose._agent import (
+from goose._internal.agent import (
     Agent,
     IAgentLogger,
     SystemMessage,
     UserMessage,
 )
-from goose._conversation import Conversation
-from goose._result import Result
+from goose._internal.conversation import Conversation
+from goose._internal.result import Result
 from goose.errors import Honk
 
 if TYPE_CHECKING:
-    from goose._task import Task
+    from goose._internal.task import Task
 
 
 @dataclass
@@ -97,23 +97,17 @@ class FlowRun:
         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)
-                )
+                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]:
-        if (
-            existing_node_state := self._node_states.get((task.name, index))
-        ) is not None:
+        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)
         else:
             return NodeState[task.result_type](
                 task_name=task.name,
                 index=index,
-                conversation=Conversation[task.result_type](
-                    user_messages=[], result_messages=[]
-                ),
+                conversation=Conversation[task.result_type](user_messages=[], result_messages=[]),
                 last_hash=0,
             )
 
@@ -143,9 +137,7 @@ class FlowRun:
         self._last_requested_indices = {}
         self._flow_name = flow_name
         self._id = run_id
-        self._agent = Agent(
-            flow_name=self.flow_name, run_id=self.id, logger=agent_logger
-        )
+        self._agent = Agent(flow_name=self.flow_name, run_id=self.id, logger=agent_logger)
 
     def end(self) -> None:
         self._last_requested_indices = {}
@@ -177,9 +169,7 @@ class FlowRun:
         return flow_run
 
 
-_current_flow_run: ContextVar[FlowRun | None] = ContextVar(
-    "current_flow_run", default=None
-)
+_current_flow_run: ContextVar[FlowRun | None] = ContextVar("current_flow_run", default=None)
 
 
 def get_current_flow_run() -> FlowRun | None:

goose_py-0.6.0/goose/_store.py → goose_py-0.7.1/goose/_internal/store.py

@@ -2,8 +2,8 @@ from __future__ import annotations
 
 from typing import Protocol
 
-from goose._flow import FlowRun
-from goose._state import FlowRunState
+from goose._internal.flow import FlowRun
+from goose._internal.state import FlowRunState
 
 
 class IFlowRunStore(Protocol):

goose_py-0.6.0/goose/_task.py → goose_py-0.7.1/goose/_internal/task.py

@@ -1,11 +1,12 @@
-from typing import Awaitable, Callable, overload
-
-from goose._agent import Agent, GeminiModel, SystemMessage, UserMessage
-from goose._conversation import Conversation
-from goose._result import Result, TextResult
-from goose._state import FlowRun, NodeState, get_current_flow_run
+from collections.abc import Awaitable, Callable
+from typing import overload
+
+from goose._internal.agent import Agent, GeminiModel, SystemMessage, UserMessage
+from goose._internal.conversation import Conversation
+from goose._internal.result import Result, TextResult
+from goose._internal.state import FlowRun, NodeState, get_current_flow_run
+from goose._internal.types.agent import AssistantMessage
 from goose.errors import Honk
-from goose.types.agent import AssistantMessage
 
 
 class Task[**P, R: Result]:
@@ -33,9 +34,7 @@ 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[R], *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)
@@ -58,9 +57,7 @@ class Task[**P, R: Result]:
             node_state.set_context(context=context)
         node_state.add_user_message(message=user_message)
 
-        result = await self.__adapt(
-            conversation=node_state.conversation, agent=flow_run.agent
-        )
+        result = await self.__adapt(conversation=node_state.conversation, agent=flow_run.agent)
         node_state.add_result(result=result)
         flow_run.add_node_state(node_state)
 
@@ -97,11 +94,7 @@ class Task[**P, R: Result]:
 
     def __hash_task_call(self, *args: P.args, **kwargs: P.kwargs) -> int:
         try:
-            to_hash = str(
-                tuple(args)
-                + tuple(kwargs.values())
-                + (self._generator.__code__, self._adapter_model)
-            )
+            to_hash = str(tuple(args) + tuple(kwargs.values()) + (self._generator.__code__, self._adapter_model))
             return hash(to_hash)
         except TypeError:
             raise Honk(f"Unhashable argument to task {self.name}: {args} {kwargs}")

goose_py-0.7.1/goose/agent.py

@@ -0,0 +1,28 @@
+from goose._internal.agent import AgentResponse, IAgentLogger
+from goose._internal.types.agent import (
+    AssistantMessage,
+    GeminiModel,
+    LLMMediaMessagePart,
+    LLMMessage,
+    LLMTextMessagePart,
+    MediaMessagePart,
+    SystemMessage,
+    TextMessagePart,
+    UserMediaContentType,
+    UserMessage,
+)
+
+__all__ = [
+    "AgentResponse",
+    "IAgentLogger",
+    "AssistantMessage",
+    "GeminiModel",
+    "LLMMediaMessagePart",
+    "LLMMessage",
+    "LLMTextMessagePart",
+    "MediaMessagePart",
+    "SystemMessage",
+    "TextMessagePart",
+    "UserMediaContentType",
+    "UserMessage",
+]

goose_py-0.7.1/goose/flow.py

@@ -0,0 +1,3 @@
+from goose._internal.flow import Flow
+
+__all__ = ["Flow"]

goose_py-0.7.1/goose/runs.py

@@ -0,0 +1,4 @@
+from goose._internal.state import FlowRun, FlowRunState
+from goose._internal.store import IFlowRunStore
+
+__all__ = ["FlowRun", "FlowRunState", "IFlowRunStore"]

goose_py-0.7.1/pyproject.toml

@@ -0,0 +1,81 @@
+[project]
+name = "goose-py"
+version = "0.7.1"
+description = "A tool for AI workflows based on human-computer collaboration and structured output."
+readme = "README.md"
+authors = [
+    { name = "Nash Taylor", email = "nash@chelle.ai" },
+    { name = "Joshua Cook", email = "joshua@chelle.ai" },
+    { name = "Michael Sankur", email = "michael@chelle.ai" },
+]
+requires-python = ">=3.12"
+dependencies = [
+    "jsonpath-ng>=1.7.0",
+    "litellm>=1.56.5",
+    "pydantic>=2.8.2",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["goose"]
+
+[dependency-groups]
+dev = [
+    "pyright>=1.1.393",
+    "pytest>=8.3.4",
+    "pytest-asyncio>=0.25.3",
+    "pytest-mock>=3.14.0",
+    "ruff>=0.9.4",
+]
+
+
+[tool.ruff]
+exclude = [
+    "__init__.py",
+    ".venv",
+    "**/.venv",
+    "notebooks",
+    ".stubs",
+]
+force-exclude = true
+line-length = 120
+
+[tool.ruff.lint]
+select = [ "E", "F", "I", "UP" ]
+ignore = [ "E501" ]
+
+
+[tool.ruff.lint.isort]
+known-first-party = [
+    "goose",
+]
+
+[tool.ruff.lint.flake8-tidy-imports]
+ban-relative-imports = "all"
+
+
+[tool.pyright]
+pythonVersion = "3.12"
+typeCheckingMode = "strict"
+reportMissingModuleSource = false
+useLibraryCodeForTypes = false
+reportUnknownMemberType = false
+reportUnknownVariableType = false
+stubPath = ".stubs"
+venvPath = "."
+venv = ".venv"
+exclude = ["goose/__init__.py", ".venv", "notebooks"]
+
+[tool.pytest.ini_options]
+filterwarnings = [
+    "ignore::DeprecationWarning",
+    "ignore::SyntaxWarning",
+    "ignore::UserWarning",
+]
+addopts = "-v"
+testpaths = ["tests"]
+pythonpath = ["."]
+python_files = "test_*.py"

goose_py-0.7.1/tests/conftest.py

@@ -0,0 +1,10 @@
+import asyncio
+
+import pytest
+
+
+@pytest.fixture(scope="session")
+def event_loop():
+    loop = asyncio.get_event_loop()
+    yield loop
+    loop.close()

goose_py-0.7.1/tests/test_agent.py

@@ -0,0 +1,66 @@
+from unittest.mock import Mock
+
+import pytest
+from pytest_mock import MockerFixture
+
+from goose import TextResult, flow, task
+from goose._internal.agent import Agent, AgentResponse, IAgentLogger
+from goose.agent import GeminiModel, TextMessagePart, UserMessage
+
+
+class MockLiteLLMResponse:
+    def __init__(self, *, response: str, prompt_tokens: int, completion_tokens: int) -> None:
+        self.choices = [Mock(message=Mock(content=response))]
+        self.usage = Mock(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens)
+
+
+@pytest.fixture
+def mock_litellm(mocker: MockerFixture) -> Mock:
+    return mocker.patch(
+        "goose._internal.agent.acompletion",
+        return_value=MockLiteLLMResponse(response="Hello", prompt_tokens=10, completion_tokens=10),
+    )
+
+
+@task
+async def use_agent(*, agent: Agent) -> TextResult:
+    return await agent(
+        messages=[UserMessage(parts=[TextMessagePart(text="Hello")])],
+        model=GeminiModel.FLASH_8B,
+        task_name="greet",
+    )
+
+
+@flow
+async def agent_flow(*, agent: Agent) -> None:
+    await use_agent(agent=agent)
+
+
+class CustomLogger(IAgentLogger):
+    logged_responses: list[AgentResponse[TextResult]] = []
+
+    async def __call__(self, *, response: AgentResponse[TextResult]) -> None:
+        self.logged_responses.append(response)
+
+
+@flow(agent_logger=CustomLogger())
+async def agent_flow_with_custom_logger(*, agent: Agent) -> None:
+    await use_agent(agent=agent)
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("mock_litellm")
+async def test_agent() -> None:
+    async with agent_flow.start_run(run_id="1") as run:
+        await agent_flow.generate(agent=run.agent)
+
+    assert run.get(task=use_agent).result.text == "Hello"
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("mock_litellm")
+async def test_agent_custom_logger() -> None:
+    async with agent_flow_with_custom_logger.start_run(run_id="1") as run:
+        await agent_flow_with_custom_logger.generate(agent=run.agent)
+
+    assert len(CustomLogger.logged_responses) == 1

goose_py-0.7.1/tests/test_complex_flow_arguments.py

@@ -0,0 +1,22 @@
+import pytest
+from pydantic import BaseModel
+
+from goose import Agent, flow
+
+
+class MyMessage(BaseModel):
+    text: str
+
+
+@flow
+async def my_flow(*, message: MyMessage, agent: Agent) -> None:
+    pass
+
+
+@pytest.mark.asyncio
+async def test_my_flow() -> None:
+    async with my_flow.start_run(run_id="1") as run:
+        await my_flow.generate(message=MyMessage(text="Hello"), agent=run.agent)
+
+    async with my_flow.start_run(run_id="1") as run:
+        await my_flow.generate(message=MyMessage(text="Hello"), agent=run.agent)