ldp 0.1.0__py3-none-any.whl

ldp/agent/__init__.py

@@ -0,0 +1,33 @@
+from enum import StrEnum
+
+
+class DefaultLLMModelNames(StrEnum):
+    """Defaults for LLM models, pin exact versions for performance stability."""
+
+    OPENAI = "gpt-4o-2024-08-06"  # Cheap, fast, and decent
+
+
+# ruff: noqa: E402  # Avoid circular imports
+
+from .agent import Agent, AgentConfig
+from .agent_client import HTTPAgentClient, make_simple_agent_server
+from .dqn_agent import DQNAgent, MultipleCompletionLLMCallOp
+from .memory_agent import MemoryAgent
+from .react_agent import ReActAgent
+from .simple_agent import SimpleAgent, SimpleAgentState
+from .tree_of_thoughts_agent import TreeofThoughtsAgent
+
+__all__ = [
+    "Agent",
+    "AgentConfig",
+    "DQNAgent",
+    "DefaultLLMModelNames",
+    "HTTPAgentClient",
+    "MemoryAgent",
+    "MultipleCompletionLLMCallOp",
+    "ReActAgent",
+    "SimpleAgent",
+    "SimpleAgentState",
+    "TreeofThoughtsAgent",
+    "make_simple_agent_server",
+]

ldp/agent/agent.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+import json
+from abc import ABC, abstractmethod
+from collections.abc import Collection, Iterable, Mapping, Sequence
+from typing import Any, Generic, TypeVar
+
+import numpy as np
+from aviary.message import Message
+from aviary.tools import Tool, ToolRequestMessage
+from pydantic import BaseModel, ConfigDict, Field, JsonValue
+
+from ldp.graph.ops import Op, OpResult
+
+try:
+    # So we can skip torch objects when looking for Ops
+    import torch
+except ImportError:
+    # If torch is not available, then it won't be used in an Agent anyway
+    torch = None  # type: ignore[assignment]
+
+TAgentState = TypeVar("TAgentState")
+
+
+# A global registry of all Agent subclasses, so we can look them up by name
+_AGENT_REGISTRY: dict[str, type[Agent]] = {}
+
+
+class Agent(ABC, Generic[TAgentState]):
+    def __init_subclass__(cls, **kwargs):
+        """Ensure Ops have unique names and subclasses are in _AGENT_REGISTRY."""
+        super().__init_subclass__(**kwargs)
+
+        original_init = cls.__init__
+
+        def init_with_op_naming(self, *args, **kwargs):
+            original_init(self, *args, **kwargs)
+
+            # loop through Ops and give them proper names
+            for name, op in _find_ops(self):
+                op.set_name(name)
+
+        cls.__init__ = init_with_op_naming  # type: ignore[method-assign]
+
+        # Register the Agent subclass.
+        _AGENT_REGISTRY[cls.__name__] = cls
+
+    @abstractmethod
+    async def get_asv(
+        self, agent_state: TAgentState, obs: list[Message]
+    ) -> tuple[OpResult[ToolRequestMessage], TAgentState, float]:
+        """
+        Get new action, state, and value given state and observation messages.
+
+        NOTE: the method's name has action listed before state to help you realize it's
+        a new state.
+
+        Args:
+            agent_state: Optional current agent state, pass None if irrelevant.
+                This can be something like agent memory.
+            obs: Most recent list of observation messages from the environment's steps.
+                If more observations than the most recent list are necessary, track them
+                in the agent state.
+
+        Returns:
+            Three-tuple of new action, new agent state, and estimated value. The
+                agent_state is returned as a copy so that you can safely mutate it
+                without affecting the original. The estimated value is the agent's
+                estimate of the future rewards given the input state and observations,
+                and is used for RL training.  If estimated value doesn't matter, just
+                return 0. The value could also come from a Q-value evaluated at the
+                action chosen by the agent.
+        """
+
+    @abstractmethod
+    async def init_state(self, tools: list[Tool]) -> TAgentState:
+        """Initializes the first agent state with the provided tools."""
+
+    def named_ops(self) -> Iterable[tuple[str, Op]]:
+        """Analogous to torch.nn.Module.named_parameters()."""
+        return _find_ops(self)
+
+
+class AgentConfig(BaseModel):
+    """Configuration for specifying the type of agent i.e. the subclass of Agent above."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    agent_type: str = Field(
+        description="The type of agent to be used. "
+        "This should be a subclass of Agent above.",
+    )
+    agent_kwargs: dict[str, JsonValue] = Field(
+        default_factory=dict,
+        description="Keyword arguments to pass to the agent's constructor.",
+    )
+
+    def construct_agent(self) -> Agent:
+        return _AGENT_REGISTRY[self.agent_type](**self.agent_kwargs)
+
+    def __hash__(self) -> int:
+        return hash(self.agent_type + json.dumps(self.agent_kwargs, sort_keys=True))
+
+
+def _find_ops(  # noqa: C901
+    root: object, root_name: str = "", visited: set[int] | None = None
+) -> Iterable[tuple[str, Op]]:
+    """Recursive function to find children that are Ops and the attr chain to reach them.
+
+    E.g. if root.module.op is an Op, then we will yield ("module.op", root.module.op).
+    These are not fully qualified names, but more like "locally qualified names". In the above
+    example, "root." + "module.op" is the fully qualified name.
+    This is an internal function - Agent.named_ops() should usually suffice.
+
+    Args:
+        root: Any object that might hold Ops.
+        root_name: The name of the root object. Defaults to empty string and is passed as an arg to
+            make this method recursive.
+        visited: a set of visited node IDs to avoid loops. Defaults to None.
+
+    Yields:
+        Two-tuple of (locally qualified name, Op) pairs
+    """
+    # Recursive function to find children that are Ops and the
+    # attribute chain to reach them.
+    if visited is None:
+        visited = set()
+
+    if isinstance(root, Op):
+        yield root_name, root
+        # Assume an Op may not have sub-Ops. I think this is sound, since
+        # we wouldn't be tracking the compute graph properly if it did.
+        return
+
+    if "__pydantic_parent_namespace__" in root_name:
+        # Skip Pydantic internals
+        return
+
+    # Don't recurse into PyTorch objects because they won't contain Ops
+    if torch is not None and (  # type: ignore[redundant-expr]
+        isinstance(root, torch.Tensor | torch.nn.Module)
+    ):
+        return
+
+    # Similarly for numpy
+    if isinstance(root, np.ndarray):
+        return
+
+    # loop through 3 types of containers: dicts, collections, and objects
+    if isinstance(root, Mapping):
+        named_attrs: Any = root.items()
+    elif isinstance(root, Sequence | Collection) and not isinstance(root, str | bytes):
+        named_attrs = enumerate(root)
+    elif hasattr(root, "__dict__"):
+        # object?
+        named_attrs = root.__dict__.items()
+    else:
+        # couldn't descend
+        return
+
+    for k, v in named_attrs:
+        id_v = id(v)
+        if id_v not in visited:
+            # only visit each object once - avoid loops, etc.
+            visited.add(id_v)
+            if root_name:
+                k = f"{root_name}.{k}"
+            yield from _find_ops(v, root_name=k, visited=visited)

ldp/agent/agent_client.py

@@ -0,0 +1,147 @@
+import os
+import secrets
+from typing import TYPE_CHECKING, Annotated, TypeVar
+
+import httpx
+from aviary.message import Message
+from aviary.tools import Messages, Tool, ToolRequestMessage, ToolsAdapter
+from pydantic import BaseModel
+
+from ldp.graph.op_utils import get_training_mode
+from ldp.graph.ops import OpResult
+
+from .agent import Agent
+from .simple_agent import SimpleAgentState
+
+if TYPE_CHECKING:
+    from fastapi import FastAPI
+
+TSerializableAgentState = TypeVar("TSerializableAgentState", bound=BaseModel)
+
+
+class HTTPAgentClient(Agent[TSerializableAgentState]):
+    """Interact with an Agent running in a server via POST requests."""
+
+    def __init__(
+        self,
+        agent_state_type: type[TSerializableAgentState],
+        server_url: str,
+        request_headers: httpx._types.HeaderTypes | None = None,
+        request_timeout: float | None = None,
+    ):
+        super().__init__()
+        self._agent_state_type = agent_state_type
+        self._request_url = server_url
+        self._request_headers = request_headers
+        self._request_timeout = request_timeout
+
+    async def get_asv(
+        self,
+        agent_state: TSerializableAgentState,
+        obs: list[Message],
+    ) -> tuple[OpResult[ToolRequestMessage], TSerializableAgentState, float]:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self._request_url}/get_asv",
+                json={
+                    "agent_state": agent_state.model_dump(),
+                    "obs": [m.model_dump() for m in obs],
+                    "training": get_training_mode(),
+                },
+                headers=self._request_headers,
+                timeout=self._request_timeout,
+            )
+            response.raise_for_status()
+            response_data = response.json()
+            return (
+                OpResult[ToolRequestMessage](**response_data[0]),
+                self._agent_state_type(**response_data[1]),
+                response_data[2],
+            )
+
+    async def init_state(self, tools: list[Tool]) -> TSerializableAgentState:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                f"{self._request_url}/init_state",
+                json=ToolsAdapter.dump_python(tools),
+                headers=self._request_headers,
+                timeout=self._request_timeout,
+            )
+            response.raise_for_status()
+            return self._agent_state_type(**response.json())
+
+
+def make_simple_agent_server(
+    agent: Agent[SimpleAgentState], render_docs: bool = False
+) -> "FastAPI":
+    """
+    Make a FastAPI app designed to work with the above HTTPAgentClient.
+
+    Here's how this works:
+    1. There is an entity orchestrating an Agent's interactions with an Environment.
+       A simple example of this is an integration test that sequentially calls
+       Agent.get_asv and Environment.step.
+    2. That entity is given the above HTTPAgentClient. Any Agent.init_state or
+       Agent.get_asv calls the orchestration entity makes are actually
+       POST requests under the hood. The agent's "brains" aren't local.
+    3. Remotely, this server code is running, and is where the actual Agent logic lives.
+       An example of this is a remote server containing GPU(s).
+    """
+    try:
+        from fastapi import Body, Depends, FastAPI, HTTPException, status
+        from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+    except ModuleNotFoundError as exc:
+        raise ImportError(
+            "Please install aviary with the 'server' extra like so:"
+            " `pip install aviary[server]`."
+        ) from exc
+
+    asgi_app = FastAPI(
+        title=f"aviary.Agent {type(agent).__name__}",
+        description="Serve inference endpoints for an aviary.Agent with a SimpleAgentState",
+        # Only render Swagger docs if local since we don't have a login here
+        docs_url="/docs" if render_docs else None,
+        redoc_url="/redoc" if render_docs else None,
+    )
+    auth_scheme = HTTPBearer()
+
+    async def validate_token(
+        token: Annotated[HTTPAuthorizationCredentials, Depends(auth_scheme)],
+    ) -> HTTPAuthorizationCredentials:
+        # NOTE: don't use os.environ.get() to avoid possible empty string matches, and
+        # to have clearer server failures if the AUTH_TOKEN env var isn't present
+        if not secrets.compare_digest(token.credentials, os.environ["AUTH_TOKEN"]):
+            raise HTTPException(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                detail="Incorrect bearer token",
+                headers={"WWW-Authenticate": "Bearer"},
+            )
+        return token
+
+    @asgi_app.get("/info")
+    def info(
+        _: Annotated[HTTPAuthorizationCredentials, Depends(validate_token)],
+    ) -> dict[str, str]:
+        """Get agent metadata, useful for debugging."""
+        return {"agent_type": type(agent).__name__}
+
+    @asgi_app.post("/get_asv")
+    async def get_asv(
+        agent_state: SimpleAgentState,
+        obs: Messages,
+        _: Annotated[HTTPAuthorizationCredentials, Depends(validate_token)],
+        training: Annotated[bool, Body()] = True,
+    ) -> tuple[OpResult[ToolRequestMessage], SimpleAgentState, float]:
+        if training:
+            raise NotImplementedError("Training is not yet supported.")
+        action, agent_state, vhat = await agent.get_asv(agent_state, obs)
+        return action, agent_state, vhat
+
+    @asgi_app.post("/init_state")
+    async def init_state(
+        tools: list[Tool],
+        _: Annotated[HTTPAuthorizationCredentials, Depends(validate_token)],
+    ) -> SimpleAgentState:
+        return await agent.init_state(tools)
+
+    return asgi_app

ldp/agent/dqn_agent.py

@@ -0,0 +1,154 @@
+import asyncio
+from typing import Any, cast
+
+from aviary.message import Message
+from aviary.tools import Tool, ToolRequestMessage
+from pydantic import BaseModel, ConfigDict, Field
+
+from ldp.graph.common_ops import FxnOp, LLMCallOp
+from ldp.graph.modules import DQNOp, DQNPolicyModule
+from ldp.graph.op_utils import compute_graph, get_call_id
+from ldp.graph.ops import OpResult
+from ldp.llms import MultipleCompletionLLMModel, prepend_sys
+
+from . import DefaultLLMModelNames
+from .agent import Agent
+from .simple_agent import SimpleAgentState
+
+
+class MultipleCompletionLLMCallOp(LLMCallOp):
+    """Like LLMCallOp, but samples multiple completions (n>1) for a given input."""
+
+    async def forward(  # type: ignore[override]
+        self,
+        config: dict,
+        msgs: list[Message],
+        tools: list[Tool] | None = None,
+        tool_choice: Tool
+        | str
+        | None = MultipleCompletionLLMModel.TOOL_CHOICE_REQUIRED,
+    ) -> list[Message]:
+        model = MultipleCompletionLLMModel(config=config)
+
+        results = await model.call(messages=msgs, tools=tools, tool_choice=tool_choice)
+        if not results:
+            raise ValueError("No completions returned from the model.")
+
+        # All completions have the same config, so check the first one
+        temperature: float = (results[0].config or {}).get("temperature", 1.0)
+
+        # Compute a Monte Carlo estimate of the logprob of this sequence at the given temperature.
+        logprobs = await asyncio.gather(*[
+            self.compute_logprob(
+                raw_log_p=result.logprob,
+                temperature=temperature,
+                model=model,
+                messages=msgs,
+                tools=tools,
+                tool_choice=tool_choice,
+            )
+            for result in results
+        ])
+
+        call_id = get_call_id()
+        self.ctx.update(call_id, "results", results)
+        # This is the logprob of this sequence according to the raw model, without
+        # any temperature/top-p distribution shaping.
+        self.ctx.update(call_id, "raw_logprobs", [result.logprob for result in results])
+
+        self.ctx.update(call_id, "temperature", temperature)
+        self.ctx.update(call_id, "logprob", logprobs)
+
+        return [cast(list[Message], result.messages)[0] for result in results]
+
+    # set the return type to list[Message], not LLMCallOp's inherited Message
+    async def __call__(  # type: ignore[override]
+        self, *args, **kwargs
+    ) -> OpResult[list[Message]]:
+        return await super().__call__(*args, **kwargs)  # type: ignore[return-value]
+
+    async def compute_logprob(
+        self,
+        raw_log_p: float | None,
+        temperature: float,
+        model: MultipleCompletionLLMModel,
+        **model_kwargs,
+    ) -> float | None:
+        """This method computes a Monte Carlo estimate of logprob for a given temperature."""
+        return None  # TODO: finish implementing this using n>1
+
+
+class DQNAgent(BaseModel, Agent[SimpleAgentState]):
+    """An agent that trains a state-action value function Q(s,a) to select actions.
+
+    This a modification of traditional DQNs [1]. When using a vanilla DQN, the action
+    space is assumed to be enumerable, enabling greedy action selection. When sampling
+    from an LLM, the action space is large. So instead, we sample from the LLM and use
+    the Q network to score the sampled actions. We are then greedy among the sampled
+    actions.
+
+    [1] https://arxiv.org/abs/1312.5602
+    """
+
+    # Not frozen because we want to change num_actions_to_sample
+    model_config = ConfigDict(frozen=False, arbitrary_types_allowed=True)
+
+    num_actions_to_sample: int = Field(
+        default=1,
+        description="Number of actions to sample from the LLM. "
+        "If >1, the Q function will be used to select the "
+        "highest-scoring action.",
+    )
+
+    llm_model: dict[str, Any] = Field(
+        default={"model": DefaultLLMModelNames.OPENAI.value, "temperature": 1.0},
+        description="Model configuration (not trained). "
+        "Setting a high default temperature to encourage exploration.",
+    )
+    sys_prompt: str = Field(
+        default="Using tools, complete the given task.",
+        description="System prompt. Trainable",
+    )
+
+    def __init__(
+        self,
+        dqn: DQNOp | None = None,
+        epsilon: float = 0.0,
+        actions_only: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self._prepend = FxnOp(prepend_sys)
+        self._llm_call = MultipleCompletionLLMCallOp()
+        self._action_splitter = FxnOp[Message](lambda x, i: x[i])
+        self._dqn_policy = DQNPolicyModule[ToolRequestMessage](
+            dqn=dqn, epsilon=epsilon, actions_only=actions_only
+        )
+
+    async def init_state(self, tools: list[Tool]) -> SimpleAgentState:
+        return SimpleAgentState(tools=tools)
+
+    @compute_graph()
+    async def get_asv(
+        self, agent_state: SimpleAgentState, obs: list[Message]
+    ) -> tuple[OpResult[ToolRequestMessage], SimpleAgentState, float]:
+        new_state = agent_state.get_next_state(obs)
+
+        msgs = await self._prepend(new_state.messages, sys_content=self.sys_prompt)
+        sampled_actions: OpResult[list[Message]] = await self._llm_call(
+            # Override config's n. Also make sure caching=False, since we always want
+            # stochastic samples from the LLM for the DQN.
+            self.llm_model | {"n": self.num_actions_to_sample, "caching": False},
+            msgs=msgs,
+            tools=agent_state.tools,
+        )
+        split_actions: list[OpResult[Message]] = await asyncio.gather(*[
+            self._action_splitter(sampled_actions, i)
+            for i in range(self.num_actions_to_sample)
+        ])
+
+        best_q, best_action = await self._dqn_policy(msgs, *split_actions)  # type: ignore[arg-type]
+        new_state.messages = [*new_state.messages, best_action.value]
+
+        return best_action, new_state, best_q

ldp/agent/memory_agent.py

@@ -0,0 +1,100 @@
+"""This module defines the MemoryAgent class, which extends a base agent model with memory.
+
+capabilities. The MemoryAgent can pick and invoke tools based on the stored and retrieved
+memories, formatted using specified prompts. A memory is typically a set of previous trajectories
+"""
+
+from typing import cast
+
+from aviary.message import Message
+from aviary.tools import ToolRequestMessage
+from pydantic import ConfigDict, Field
+
+from ldp.graph.common_ops import FxnOp, MemoryOp, PromptOp
+from ldp.graph.memory import Memory, MemoryModel
+from ldp.graph.op_utils import compute_graph
+from ldp.graph.ops import OpResult
+from ldp.llms.prompts import indent_xml
+
+from .simple_agent import SimpleAgent, SimpleAgentState
+
+
+class MemoryAgent(SimpleAgent):
+    """
+    Simple agent that can pick and invoke tools with memory.
+
+    NOTE: the MemoryAgent does not maintain an explicit value estimate,
+    it simply supplies previous trajectories via the prompt.
+    As such, the value estimate vhat will always be zero.
+    """
+
+    prompt: str = Field(
+        default=(
+            "<episode-memories>\n<description>\n"
+            "These are relevant memories from previous attempts at similar tasks, "
+            "along with the action taken and the discounted cumulative reward from that action. "
+            "A negative reward is failure, a positive reward is success.\n"
+            "</description>{memories}</episode-memories>\n\n"
+            "Considering the memories, choose the next action."
+        ),
+        description="Prompt that includes the memories.",
+    )
+    memory_prompt: str = Field(
+        default="<memory><obs>{input}</obs><action>{output}</action><reward>{value}</reward></memory>",
+        description="Prompt for formatting an individual memory. "
+        "Use XML instead of JSON to avoid potential escaping issues.",
+    )
+    num_memories: int = Field(
+        default=MemoryModel.DEFAULT_MEMORY_MATCHES,
+        description="Number of memories to retrieve from MemoryOp",
+    )
+    # Freeze to ensure the only mutation happens in either the agent state (which is
+    # passed around) or in the internal Ops
+    model_config = ConfigDict(frozen=True)
+
+    @staticmethod
+    def _parse_memory(prompt: str, memories: list[Memory]) -> str:
+        return indent_xml(
+            "\n".join([
+                prompt.format(**m.model_dump(exclude={"call_id"})) for m in memories
+            ])
+        )
+
+    @staticmethod
+    def _package_messages(
+        msgs: list[Message], memory_prompt: str, use_memories: bool
+    ) -> list[Message]:
+        if use_memories:
+            return [*msgs, Message(content=memory_prompt)]
+        return msgs
+
+    def __init__(self, memory_model: MemoryModel | None = None, **kwargs):
+        super().__init__(**kwargs)
+        self._memory_op = MemoryOp(memory_model)
+        self._format_memory_op = FxnOp(self._parse_memory)
+        self._prompt_op = PromptOp(self.prompt)
+        self._package_op = FxnOp(self._package_messages)
+
+    @compute_graph()
+    async def get_asv(
+        self, agent_state: SimpleAgentState, obs: list[Message]
+    ) -> tuple[OpResult[ToolRequestMessage], SimpleAgentState, float]:
+        next_state = agent_state.get_next_state(obs)
+
+        query = "\n\n".join([str(m) for m in next_state.messages if m.role != "system"])
+        memories = await self._memory_op(query, matches=self.num_memories)
+        packaged_messages = await self._package_op(
+            next_state.messages,
+            memory_prompt=await self._prompt_op(
+                memories=await self._format_memory_op(self.memory_prompt, memories)
+            ),
+            use_memories=bool(memories.value),
+        )
+        result = cast(
+            OpResult[ToolRequestMessage],
+            await self._llm_call_op(
+                await self._config_op(), msgs=packaged_messages, tools=next_state.tools
+            ),
+        )
+        next_state.messages = [*next_state.messages, result.value]
+        return result, next_state, 0.0