goose-py 0.1.0__tar.gz

goose_py-0.1.0/PKG-INFO

@@ -0,0 +1,31 @@
+Metadata-Version: 2.1
+Name: goose-py
+Version: 0.1.0
+Summary: A tool for AI workflows based on human-computer collaboration and structured output.
+Home-page: https://github.com/chelle-ai/goose
+Keywords: ai,yaml,configuration,llm
+Author: Nash Taylor
+Author-email: nash@chelle.ai
+Requires-Python: >=3.12,<4.0
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Provides-Extra: test
+Requires-Dist: ipykernel ; extra == "test"
+Requires-Dist: jsonpath-ng (>=1.7.0,<2.0.0)
+Requires-Dist: litellm (>=1.56.5,<2.0.0)
+Requires-Dist: pydantic (>=2.8.2,<3.0.0)
+Requires-Dist: pytest (<8) ; extra == "test"
+Requires-Dist: pytest-asyncio ; extra == "test"
+Requires-Dist: pytest-mock ; extra == "test"
+Project-URL: Documentation, https://github.com/chelle-ai/goose
+Project-URL: Repository, https://github.com/chelle-ai/goose
+Description-Content-Type: text/markdown
+
+# Vers
+
+Docs to come.
+

goose_py-0.1.0/README.md

@@ -0,0 +1,3 @@
+# Vers
+
+Docs to come.

goose_py-0.1.0/goose/agent.py

@@ -0,0 +1,74 @@
+import logging
+import uuid
+from datetime import datetime
+from typing import Any, Callable
+
+from litellm import acompletion
+from pydantic import BaseModel
+
+from goose.types import (
+    AgentResponse,
+    AssistantMessage,
+    GeminiModel,
+    SystemMessage,
+    UserMessage,
+)
+
+
+class Agent:
+    def __init__(
+        self,
+        *,
+        flow_name: str,
+        logger: Callable[[AgentResponse[Any]], None] | None = None,
+    ) -> None:
+        self.flow_name = flow_name
+        self.logger = logger or logging.info
+
+    async def __call__[R: BaseModel](
+        self,
+        *,
+        messages: list[UserMessage | AssistantMessage],
+        model: GeminiModel,
+        response_model: type[R],
+        task_name: str,
+        system: SystemMessage | None = None,
+    ) -> R:
+        start_time = datetime.now()
+        rendered_messages = [message.render() for message in messages]
+        if system is not None:
+            rendered_messages.insert(0, system.render())
+
+        response = await acompletion(
+            model=model.value,
+            messages=rendered_messages,
+            response_format={
+                "type": "json_object",
+                "response_schema": response_model.model_json_schema(),
+                "enforce_validation": True,
+            },
+        )
+
+        if len(response.choices) == 0:
+            raise RuntimeError("No content returned from LLM call.")
+
+        parsed_response = response_model.model_validate_json(
+            response.choices[0].message.content
+        )
+        end_time = datetime.now()
+        agent_response = AgentResponse(
+            response=parsed_response,
+            id=str(uuid.uuid4()),
+            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,
+        )
+
+        self.logger(agent_response)
+        return agent_response.response

goose_py-0.1.0/goose/conversation.py

@@ -0,0 +1,45 @@
+from pydantic import BaseModel
+
+from goose.types import UserMessage
+
+
+class ConversationState[R: BaseModel](BaseModel):
+    user_messages: list[UserMessage]
+    results: list[R]
+
+
+class Conversation[R: BaseModel]:
+    def __init__(
+        self,
+        *,
+        user_messages: list[UserMessage] | None = None,
+        results: list[R] | None = None,
+    ) -> None:
+        self.user_messages = user_messages or []
+        self.results = results or []
+
+    @classmethod
+    def load(cls, *, state: ConversationState[R]) -> "Conversation[R]":
+        return cls(user_messages=state.user_messages, results=state.results)
+
+    @property
+    def current_result(self) -> R:
+        if len(self.results) == 0:
+            raise RuntimeError("No results in conversation")
+
+        return self.results[-1]
+
+    def add_message(self, *, message: UserMessage) -> None:
+        self.user_messages.append(message)
+
+    def add_result(self, *, result: R) -> None:
+        self.results.append(result)
+
+    def replace_last_result(self, *, result: R) -> None:
+        if len(self.results) == 0:
+            self.results.append(result)
+        else:
+            self.results[-1] = result
+
+    def dump(self) -> ConversationState[R]:
+        return ConversationState(user_messages=self.user_messages, results=self.results)

goose_py-0.1.0/goose/core.py

@@ -0,0 +1,294 @@
+import asyncio
+import contextvars
+import inspect
+from collections import defaultdict
+from types import TracebackType
+from typing import Any, Awaitable, Callable, Protocol, Self, overload
+
+from graphlib import TopologicalSorter
+from pydantic import BaseModel
+
+from goose.agent import Agent
+from goose.conversation import Conversation, ConversationState
+from goose.regenerator import default_regenerator
+from goose.types import AgentResponse, UserMessage
+
+
+class NodeState[R: BaseModel](BaseModel):
+    name: str
+    conversation: ConversationState[R]
+
+    @property
+    def result(self) -> R:
+        return self.conversation.results[-1]
+
+
+class FlowState(BaseModel):
+    nodes: list[NodeState[BaseModel]]
+
+
+class NoResult:
+    pass
+
+
+class IRegenerator[R: BaseModel](Protocol):
+    async def __call__(self, *, result: R, conversation: Conversation[R]) -> R: ...
+
+
+class Task[**P, R: BaseModel]:
+    def __init__(
+        self, generator: Callable[P, Awaitable[R]], /, *, retries: int = 0
+    ) -> None:
+        self.retries = retries
+        self._generator = generator
+        self._regenerator: IRegenerator[R] = default_regenerator
+        self._signature = inspect.signature(generator)
+        self.__validate_fn()
+
+    @property
+    def result_type(self) -> type[R]:
+        return_type = self._generator.__annotations__.get("return")
+        if return_type is None:
+            raise TypeError("Task must have a return type annotation")
+
+        return return_type
+
+    @property
+    def name(self) -> str:
+        return self._generator.__name__
+
+    def regenerator(self, regenerator: IRegenerator[R], /) -> Self:
+        self._regenerator = regenerator
+        return self
+
+    async def generate(self, *args: P.args, **kwargs: P.kwargs) -> R:
+        return await self._generator(*args, **kwargs)
+
+    async def regenerate(self, *, result: R, conversation: Conversation[R]) -> R:
+        return await self._regenerator(result=result, conversation=conversation)
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> "Node[R]":
+        arguments = self._signature.bind(*args, **kwargs).arguments
+        return Node(task=self, arguments=arguments, result_type=self.result_type)
+
+    def __validate_fn(self) -> None:
+        if any(
+            param.kind == inspect.Parameter.POSITIONAL_ONLY
+            for param in self._signature.parameters.values()
+        ):
+            raise ValueError("Positional-only parameters are not supported in Tasks")
+
+
+class Node[R: BaseModel]:
+    def __init__(
+        self,
+        *,
+        task: Task[Any, R],
+        arguments: dict[str, Any],
+        result_type: type[R],
+        conversation: Conversation[R] | None = None,
+    ) -> None:
+        self.task = task
+        self.arguments = arguments
+        self.result_type = result_type
+        self.conversation = conversation or Conversation[R]()
+        self.name = task.name
+
+        self._result: R | NoResult = NoResult()
+        current_flow = Flow.get_current()
+        if current_flow is None:
+            raise RuntimeError("Cannot create a node without an active flow")
+        self.id = current_flow.add_node(node=self)
+
+    @property
+    def has_result(self) -> bool:
+        return not isinstance(self._result, NoResult)
+
+    @property
+    def result(self) -> R:
+        if isinstance(self._result, NoResult):
+            raise RuntimeError("Cannot access result of a node before it has run")
+        return self._result
+
+    async def generate(self) -> None:
+        self._result = await self.task.generate(**self.arguments)
+        self.conversation.replace_last_result(result=self.result)
+
+    async def regenerate(self, *, message: UserMessage) -> None:
+        self.conversation.add_message(message=message)
+        self._result = await self.task.regenerate(
+            result=self.result, conversation=self.conversation
+        )
+        self.conversation.add_result(result=self.result)
+
+    def dump_state(self) -> NodeState[R]:
+        return NodeState(name=self.name, conversation=self.conversation.dump())
+
+    def load_state(self, *, state: NodeState[R]) -> None:
+        self._result = state.result
+        self.conversation = Conversation[self.result_type].load(
+            state=state.conversation
+        )
+
+    def get_inbound_nodes(self) -> list["Node[BaseModel]"]:
+        def __find_nodes(
+            obj: Any, visited: set[int] | None = None
+        ) -> list["Node[BaseModel]"]:
+            if visited is None:
+                visited = set()
+
+            if isinstance(obj, Node):
+                return [obj]
+            elif isinstance(obj, dict):
+                return [
+                    node
+                    for value in obj.values()
+                    for node in __find_nodes(value, visited)
+                ]
+            elif isinstance(obj, list):
+                return [node for item in obj for node in __find_nodes(item, visited)]
+            elif isinstance(obj, tuple):
+                return [node for item in obj for node in __find_nodes(item, visited)]
+            elif isinstance(obj, set):
+                return [node for item in obj for node in __find_nodes(item, visited)]
+            elif hasattr(obj, "__dict__"):
+                return [
+                    node
+                    for value in obj.__dict__.values()
+                    for node in __find_nodes(value, visited)
+                ]
+            return []
+
+        return __find_nodes(self.arguments)
+
+    def __hash__(self) -> int:
+        return hash(self.id)
+
+
+class Flow:
+    _current: contextvars.ContextVar["Flow | None"] = contextvars.ContextVar(
+        "current_flow", default=None
+    )
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        agent_logger: Callable[[AgentResponse[Any]], None] | None = None,
+    ) -> None:
+        self.name = name
+        self._nodes: list[Node[BaseModel]] = []
+        self._agent = Agent(flow_name=self.name, logger=agent_logger)
+
+    @property
+    def agent(self) -> Agent:
+        return self._agent
+
+    def dump_state(self) -> FlowState:
+        return FlowState(nodes=[node.dump_state() for node in self._nodes])
+
+    def load_state(self, *, flow_state: FlowState) -> None:
+        nodes_by_name = {node.name: node for node in self._nodes}
+        for node_state in flow_state.nodes:
+            matching_node = nodes_by_name.get(node_state.name)
+            if matching_node is None:
+                raise RuntimeError(
+                    f"Node {node_state.name} from state not found in flow"
+                )
+
+            matching_node.load_state(state=node_state)
+
+    async def generate(self) -> None:
+        graph = {node: node.get_inbound_nodes() for node in self._nodes}
+        sorter = TopologicalSorter(graph)
+        sorter.prepare()
+
+        async with asyncio.TaskGroup() as task_group:
+            while sorter.is_active():
+                ready_nodes = list(sorter.get_ready())
+                if ready_nodes:
+                    for node in ready_nodes:
+                        task_group.create_task(node.generate())
+                    sorter.done(*ready_nodes)
+                else:
+                    await asyncio.sleep(0)
+
+    async def regenerate(self, *, target: Node[Any], message: UserMessage) -> None:
+        if not target.has_result:
+            raise RuntimeError("Cannot regenerate a node without a result")
+
+        await target.regenerate(message=message)
+
+        # regenerate all downstream nodes
+        full_graph = {node: node.get_inbound_nodes() for node in self._nodes}
+        reversed_graph: dict[Node[BaseModel], set[Node[BaseModel]]] = defaultdict(set)
+        for node, inbound_nodes in full_graph.items():
+            for inbound_node in inbound_nodes:
+                reversed_graph[inbound_node].add(node)
+
+        subgraph: dict[Node[BaseModel], set[Node[BaseModel]]] = defaultdict(set)
+        queue: list[Node[BaseModel]] = [target]
+
+        while len(queue) > 0:
+            node = queue.pop(0)
+            outbound_nodes = reversed_graph[node]
+            for outbound_node in outbound_nodes:
+                subgraph[outbound_node].add(node)
+                if outbound_node not in subgraph:
+                    queue.append(outbound_node)
+
+        if len(subgraph) > 0:
+            sorter = TopologicalSorter(subgraph)
+            sorter.prepare()
+
+            async with asyncio.TaskGroup() as task_group:
+                while sorter.is_active():
+                    ready_nodes = list(sorter.get_ready())
+                    if len(ready_nodes) > 0:
+                        for node in ready_nodes:
+                            if node != target:
+                                task_group.create_task(node.generate())
+                        sorter.done(*ready_nodes)
+                    else:
+                        await asyncio.sleep(0)
+
+    @classmethod
+    def get_current(cls) -> "Flow | None":
+        return cls._current.get()
+
+    def add_node(self, *, node: Node[Any]) -> str:
+        existing_names = [node.name for node in self._nodes]
+        number = sum(1 for name in existing_names if name == node.name)
+        self._nodes.append(node)
+        node_id = f"{node.name}_{number}"
+        return node_id
+
+    def __enter__(self) -> Self:
+        if self._current.get() is not None:
+            raise RuntimeError(
+                "Cannot enter a new flow while another flow is already active"
+            )
+        self._current.set(self)
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self._current.set(None)
+
+
+@overload
+def task[**P, R: BaseModel](fn: Callable[P, Awaitable[R]], /) -> Task[P, R]: ...
+@overload
+def task[**P, R: BaseModel](
+    *, retries: int = 0
+) -> Callable[[Callable[P, Awaitable[R]]], Task[P, R]]: ...
+def task[**P, R: BaseModel](
+    fn: Callable[P, Awaitable[R]] | None = None, /, *, retries: int = 0
+) -> Task[P, R] | Callable[[Callable[P, Awaitable[R]]], Task[P, R]]:
+    if fn is None:
+        return lambda fn: Task(fn, retries=retries)
+    return Task(fn, retries=retries)

goose_py-0.1.0/goose/regenerator.py

@@ -0,0 +1,5 @@
+from goose.conversation import Conversation
+
+
+async def default_regenerator[R](*, result: R, conversation: Conversation[R]) -> R:
+    return result

goose_py-0.1.0/goose/types.py

@@ -0,0 +1,127 @@
+import base64
+from datetime import datetime
+from enum import StrEnum
+from typing import ClassVar, Literal, NotRequired, TypedDict
+
+from pydantic import BaseModel, computed_field
+
+
+class GeminiModel(StrEnum):
+    EXP = "gemini/gemini-exp-1121"
+    PRO = "gemini/gemini-1.5-pro"
+    FLASH = "gemini/gemini-1.5-flash"
+    FLASH_8B = "gemini/gemini-1.5-flash-8b"
+
+
+class UserMediaContentType(StrEnum):
+    JPEG = "image/jpeg"
+    PNG = "image/png"
+    WEBP = "image/webp"
+    MP3 = "audio/mpeg"
+    WAV = "audio/wav"
+
+
+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 TextMessagePart(BaseModel):
+    text: str
+
+    def render(self) -> LLMTextMessagePart:
+        return {"type": "text", "text": self.text}
+
+
+class MediaMessagePart(BaseModel):
+    content_type: UserMediaContentType
+    content: bytes
+
+    def render(self) -> LLMMediaMessagePart:
+        return {
+            "type": "image_url",
+            "image_url": f"data:{self.content_type};base64,{base64.b64encode(self.content).decode()}",
+        }
+
+
+class UserMessage(BaseModel):
+    parts: list[TextMessagePart | MediaMessagePart]
+
+    def render(self) -> LLMMessage:
+        content: LLMMessage = {
+            "role": "user",
+            "content": [part.render() for part in self.parts],
+        }
+        if any(isinstance(part, MediaMessagePart) 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):
+    text: str
+
+    def render(self) -> LLMMessage:
+        return {"role": "system", "content": [{"type": "text", "text": self.text}]}
+
+
+class AgentResponse[R: BaseModel](BaseModel):
+    INPUT_CENTS_PER_MILLION_TOKENS: ClassVar[dict[GeminiModel, float]] = {
+        GeminiModel.FLASH_8B: 30,
+        GeminiModel.FLASH: 15,
+        GeminiModel.PRO: 500,
+        GeminiModel.EXP: 0,
+    }
+    OUTPUT_CENTS_PER_MILLION_TOKENS: ClassVar[dict[GeminiModel, float]] = {
+        GeminiModel.FLASH_8B: 30,
+        GeminiModel.FLASH: 15,
+        GeminiModel.PRO: 500,
+        GeminiModel.EXP: 0,
+    }
+
+    response: R
+    id: str
+    flow_name: str
+    task_name: str
+    model: GeminiModel
+    system: SystemMessage | None = None
+    input_messages: list[UserMessage | AssistantMessage]
+    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 total_cost(self) -> float:
+        input_cost = self.INPUT_CENTS_PER_MILLION_TOKENS[self.model] * self.input_tokens
+        output_cost = (
+            self.OUTPUT_CENTS_PER_MILLION_TOKENS[self.model] * self.output_tokens
+        )
+        return input_cost + output_cost

goose_py-0.1.0/pyproject.toml

@@ -0,0 +1,62 @@
+[tool.poetry]
+name = "goose-py"
+version = "0.1.0"
+description = "A tool for AI workflows based on human-computer collaboration and structured output."
+authors = [
+    "Nash Taylor <nash@chelle.ai>",
+    "Joshua Cook <joshua@chelle.ai>",
+    "Michael Sankur <michael@chelle.ai>"
+]
+readme = "README.md"
+homepage = "https://github.com/chelle-ai/goose"
+repository = "https://github.com/chelle-ai/goose"
+documentation = "https://github.com/chelle-ai/goose"
+keywords = ["ai", "yaml", "configuration", "llm"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.12",
+]
+packages = [
+    { include = "goose" }
+]
+
+
+[tool.poetry.dependencies]
+python = "^3.12"
+
+jsonpath-ng = "^1.7.0"
+litellm = "^1.56.5"
+pydantic = "^2.8.2"
+ipykernel = { version = "*", optional = true }
+pytest = { version = "<8", optional = true }
+pytest-mock = { version = "*", optional = true }
+pytest-asyncio = { version = "*", optional = true }
+
+
+[tool.poetry.extras]
+test = ["ipykernel", "pytest", "pytest-mock", "pytest-asyncio"]
+
+
+[build-system]
+requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning>=1.0.0"]
+build-backend = "poetry_dynamic_versioning.backend"
+
+[tool.pyright]
+pythonVersion = "3.12"
+typeCheckingMode = "strict"
+reportMissingModuleSource = false
+useLibraryCodeForTypes = false
+reportImportCycles = true
+reportUnknownMemberType = false
+reportUnknownVariableType = false
+stubPath = ".stubs"
+
+[tool.pytest.ini_options]
+filterwarnings = [
+    "ignore::DeprecationWarning",
+    "ignore::SyntaxWarning",
+    "ignore::UserWarning",
+]