goose-py 0.5.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

goose/_internal/task.py

@@ -0,0 +1,136 @@
+from typing import Awaitable, Callable, 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
+
+
+class Task[**P, R: Result]:
+    def __init__(
+        self,
+        generator: Callable[P, Awaitable[R]],
+        /,
+        *,
+        retries: int = 0,
+        adapter_model: GeminiModel = GeminiModel.FLASH,
+    ) -> None:
+        self._generator = generator
+        self._retries = retries
+        self._adapter_model = adapter_model
+        self._adapter_model = adapter_model
+
+    @property
+    def result_type(self) -> type[R]:
+        result_type = self._generator.__annotations__.get("return")
+        if result_type is None:
+            raise Honk(f"Task {self.name} has no return type annotation")
+        return result_type
+
+    @property
+    def name(self) -> str:
+        return self._generator.__name__
+
+    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)
+            state.add_result(result=result, new_hash=state_hash, overwrite=True)
+            return result
+        else:
+            return state.result
+
+    async def jam(
+        self,
+        *,
+        user_message: UserMessage,
+        context: SystemMessage | None = None,
+        index: int = 0,
+    ) -> R:
+        flow_run = self.__get_current_flow_run()
+        node_state = flow_run.get(task=self, index=index)
+
+        if context is not None:
+            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
+        )
+        node_state.add_result(result=result)
+        flow_run.add_node_state(node_state)
+
+        return result
+
+    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)
+        result = await self.generate(node_state, *args, **kwargs)
+        flow_run.add_node_state(node_state)
+        return result
+
+    async def __adapt(self, *, conversation: Conversation[R], agent: Agent) -> R:
+        messages: list[UserMessage | AssistantMessage] = []
+        for message_index in range(len(conversation.user_messages)):
+            user_message = conversation.user_messages[message_index]
+            result = conversation.result_messages[message_index]
+
+            if isinstance(result, TextResult):
+                assistant_text = result.text
+            else:
+                assistant_text = result.model_dump_json()
+            assistant_message = AssistantMessage(text=assistant_text)
+            messages.append(assistant_message)
+            messages.append(user_message)
+
+        return await agent(
+            messages=messages,
+            model=self._adapter_model,
+            task_name=f"adapt--{self.name}",
+            system=conversation.context,
+            response_model=self.result_type,
+        )
+
+    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)
+            )
+            return hash(to_hash)
+        except TypeError:
+            raise Honk(f"Unhashable argument to task {self.name}: {args} {kwargs}")
+
+    def __get_current_flow_run(self) -> FlowRun:
+        run = get_current_flow_run()
+        if run is None:
+            raise Honk("No current flow run")
+        return run
+
+
+@overload
+def task[**P, R: Result](generator: Callable[P, Awaitable[R]], /) -> Task[P, R]: ...
+@overload
+def task[**P, R: Result](
+    *, retries: int = 0, adapter_model: GeminiModel = GeminiModel.FLASH
+) -> Callable[[Callable[P, Awaitable[R]]], Task[P, R]]: ...
+def task[**P, R: Result](
+    generator: Callable[P, Awaitable[R]] | None = None,
+    /,
+    *,
+    retries: int = 0,
+    adapter_model: GeminiModel = GeminiModel.FLASH,
+) -> Task[P, R] | Callable[[Callable[P, Awaitable[R]]], Task[P, R]]:
+    if generator is None:
+
+        def decorator(fn: Callable[P, Awaitable[R]]) -> Task[P, R]:
+            return Task(fn, retries=retries, adapter_model=adapter_model)
+
+        return decorator
+
+    return Task(generator, retries=retries, adapter_model=adapter_model)

goose/_internal/types/agent.py

@@ -0,0 +1,92 @@
+from enum import StrEnum
+from typing import Literal, NotRequired, TypedDict
+
+from pydantic import BaseModel
+
+
+class GeminiModel(StrEnum):
+    PRO = "vertex_ai/gemini-1.5-pro"
+    FLASH = "vertex_ai/gemini-1.5-flash"
+    FLASH_8B = "vertex_ai/gemini-1.5-flash-8b"
+
+
+class UserMediaContentType(StrEnum):
+    # 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 TextMessagePart(BaseModel):
+    text: str
+
+    def render(self) -> LLMTextMessagePart:
+        return {"type": "text", "text": self.text}
+
+
+class MediaMessagePart(BaseModel):
+    content_type: UserMediaContentType
+    content: str
+
+    def render(self) -> LLMMediaMessagePart:
+        return {
+            "type": "image_url",
+            "image_url": f"data:{self.content_type};base64,{self.content}",
+        }
+
+
+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):
+    parts: list[TextMessagePart | MediaMessagePart]
+
+    def render(self) -> LLMMessage:
+        return {
+            "role": "system",
+            "content": [part.render() for part in self.parts],
+        }

goose/agent.py

@@ -1,283 +1,28 @@
-import json
-import logging
-from datetime import datetime
-from enum import StrEnum
-from typing import Any, ClassVar, Literal, NotRequired, Protocol, TypedDict
-
-from litellm import acompletion
-from pydantic import BaseModel, computed_field
-from goose.result import Result, TextResult
-
-
-class GeminiModel(StrEnum):
-    PRO = "vertex_ai/gemini-1.5-pro"
-    FLASH = "vertex_ai/gemini-1.5-flash"
-    FLASH_8B = "vertex_ai/gemini-1.5-flash-8b"
-
-
-class UserMediaContentType(StrEnum):
-    # 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 TextMessagePart(BaseModel):
-    text: str
-
-    def render(self) -> LLMTextMessagePart:
-        return {"type": "text", "text": self.text}
-
-
-class MediaMessagePart(BaseModel):
-    content_type: UserMediaContentType
-    content: str
-
-    def render(self) -> LLMMediaMessagePart:
-        return {
-            "type": "image_url",
-            "image_url": f"data:{self.content_type};base64,{self.content}",
-        }
-
-
-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):
-    parts: list[TextMessagePart | MediaMessagePart]
-
-    def render(self) -> LLMMessage:
-        return {
-            "role": "system",
-            "content": [part.render() for part in self.parts],
-        }
-
-
-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_CENTS_PER_MILLION_TOKENS: ClassVar[dict[GeminiModel, float]] = {
-        GeminiModel.FLASH_8B: 30,
-        GeminiModel.FLASH: 15,
-        GeminiModel.PRO: 500,
-    }
-    OUTPUT_CENTS_PER_MILLION_TOKENS: ClassVar[dict[GeminiModel, float]] = {
-        GeminiModel.FLASH_8B: 30,
-        GeminiModel.FLASH: 15,
-        GeminiModel.PRO: 500,
-    }
-
-    response: R
-    run_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 input_cost(self) -> float:
-        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
-        )
-
-    @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.render()
-            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.render() 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,
-        }
-
-
-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 __call__[R: Result](
-        self,
-        *,
-        messages: list[UserMessage | AssistantMessage],
-        model: GeminiModel,
-        task_name: str,
-        response_model: type[R] = TextResult,
-        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())
-
-        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}
-            )
-        else:
-            response = await acompletion(
-                model=model.value,
-                messages=rendered_messages,
-                response_format={
-                    "type": "json_object",
-                    "response_schema": response_model.model_json_schema(),
-                    "enforce_validation": True,
-                },
-            )
-            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
+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",
+]