Use permchain agent executor, Streaming, API docs

backend/app/storage.py

@@ -2,8 +2,8 @@
 from datetime import datetime
 
 import orjson
-from langchain.schema.messages import messages_from_dict
 from langchain.utilities.redis import get_client
+from agent_executor.checkpoint import RedisCheckpoint
 from redis.client import Redis as RedisType
 
 
@@ -115,13 +115,13 @@ def list_threads(user_id: str):
 
 
 def get_thread_messages(user_id: str, thread_id: str):
-    client = _get_redis_client()
-    messages = client.lrange(thread_messages_key(user_id, thread_id), 0, -1)
+    client = RedisCheckpoint()
+    checkpoint = client.get(
+        {"configurable": {"user_id": user_id, "thread_id": thread_id}}
+    )
+    _, messages = checkpoint.get("messages", [[], []])
     return {
-        "messages": [
-            m.dict()
-            for m in messages_from_dict([orjson.loads(m) for m in messages[::-1]])
-        ],
+        "messages": [m.dict() for m in messages],
     }
 
 

backend/packages/agent-executor/agent_executor/checkpoint.py

@@ -0,0 +1,65 @@
+import os
+import pickle
+from functools import partial
+from typing import Any, Mapping, Sequence
+
+from langchain.pydantic_v1 import Field
+from langchain.schema.runnable import RunnableConfig
+from langchain.schema.runnable.utils import ConfigurableFieldSpec
+from langchain.utilities.redis import get_client
+from permchain.checkpoint.base import BaseCheckpointAdapter
+from redis.client import Redis as RedisType
+
+
+def checkpoint_key(user_id: str, thread_id: str):
+    return f"opengpts:{user_id}:thread:{thread_id}:checkpoint"
+
+
+def _dump(mapping: dict[str, Any]) -> dict:
+    return {k: pickle.dumps(v) if v is not None else None for k, v in mapping.items()}
+
+
+def _load(mapping: dict[bytes, bytes]) -> dict:
+    return {
+        k.decode(): pickle.loads(v) if v is not None else None
+        for k, v in mapping.items()
+    }
+
+
+class RedisCheckpoint(BaseCheckpointAdapter):
+    client: RedisType = Field(
+        default_factory=partial(get_client, os.environ.get("REDIS_URL"))
+    )
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @property
+    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+        return [
+            ConfigurableFieldSpec(
+                id="user_id",
+                annotation=str,
+                name="User ID",
+                description=None,
+                default=None,
+            ),
+            ConfigurableFieldSpec(
+                id="thread_id",
+                annotation=str,
+                name="Thread ID",
+                description=None,
+                default="",
+            ),
+        ]
+
+    def _hash_key(self, config: RunnableConfig) -> str:
+        return checkpoint_key(
+            config["configurable"]["user_id"], config["configurable"]["thread_id"]
+        )
+
+    def get(self, config: RunnableConfig) -> Mapping[str, Any] | None:
+        return _load(self.client.hgetall(self._hash_key(config)))
+
+    def put(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None:
+        return self.client.hmset(self._hash_key(config), _dump(checkpoint))

backend/packages/agent-executor/agent_executor/permchain.py

@@ -0,0 +1,105 @@
+import json
+
+from permchain import Channel, Pregel
+from permchain.channels import Topic
+from permchain.checkpoint.base import BaseCheckpointAdapter
+from langchain.schema.runnable import (
+    Runnable,
+    RunnableConfig,
+    RunnableLambda,
+    RunnablePassthrough,
+)
+from langchain.schema.agent import AgentAction, AgentFinish, AgentActionMessageLog
+from langchain.schema.messages import AIMessage, FunctionMessage, AnyMessage
+from langchain.tools import BaseTool
+
+
+def _create_agent_message(
+    output: AgentAction | AgentFinish
+) -> list[AnyMessage] | AnyMessage:
+    if isinstance(output, AgentAction):
+        if isinstance(output, AgentActionMessageLog):
+            output.message_log[-1].additional_kwargs["agent"] = output
+            return output.message_log
+        else:
+            return AIMessage(
+                content=output.log,
+                additional_kwargs={"agent": output},
+            )
+    else:
+        return AIMessage(
+            content=output.return_values["output"],
+            additional_kwargs={"agent": output},
+        )
+
+
+def _create_function_message(
+    agent_action: AgentAction, observation: str
+) -> FunctionMessage:
+    if not isinstance(observation, str):
+        try:
+            content = json.dumps(observation, ensure_ascii=False)
+        except Exception:
+            content = str(observation)
+    else:
+        content = observation
+    return FunctionMessage(
+        name=agent_action.tool,
+        content=content,
+    )
+
+
+def _run_tool(
+    messages: list[AnyMessage], config: RunnableConfig, *, tools: dict[str, BaseTool]
+) -> FunctionMessage:
+    action: AgentAction = messages[-1].additional_kwargs["agent"]
+    tool = tools[action.tool]
+    result = tool.invoke(action.tool_input, config)
+    return _create_function_message(action, result)
+
+
+async def _arun_tool(
+    messages: list[AnyMessage], config: RunnableConfig, *, tools: dict[str, BaseTool]
+) -> FunctionMessage:
+    action: AgentAction = messages[-1].additional_kwargs["agent"]
+    tool = tools[action.tool]
+    result = await tool.ainvoke(action.tool_input, config)
+    return _create_function_message(action, result)
+
+
+def get_agent_executor(
+    tools: list[BaseTool],
+    agent: Runnable[dict[str, list[AnyMessage]], AgentAction | AgentFinish],
+    checkpoint: BaseCheckpointAdapter,
+) -> Pregel:
+    tool_map = {tool.name: tool for tool in tools}
+    tool_lambda = RunnableLambda(_run_tool, _arun_tool).bind(tools=tool_map)
+
+    tool_chain = tool_lambda | Channel.write_to("messages")
+    agent_chain = (
+        {"messages": RunnablePassthrough()}
+        | agent
+        | _create_agent_message
+        | Channel.write_to("messages")
+    )
+
+    def route_last_message(messages: list[AnyMessage]) -> Runnable:
+        message = messages[-1]
+        if isinstance(message, AIMessage):
+            if isinstance(message.additional_kwargs.get("agent"), AgentAction):
+                # TODO if this is last step, return stop message instead
+                return tool_chain
+            elif isinstance(message.additional_kwargs.get("agent"), AgentFinish):
+                return RunnablePassthrough()
+        else:
+            return agent_chain
+
+    executor = Channel.subscribe_to("messages") | route_last_message
+
+    return Pregel(
+        chains={"executor": executor},
+        channels={"messages": Topic(AnyMessage, accumulate=True)},
+        input=["messages"],
+        output=["messages"],
+        checkpoint=checkpoint,
+    )

backend/packages/gizmo-agent/gizmo_agent/agent_types/openai.py

@@ -1,6 +1,5 @@
 import os
 
-from langchain.agents.format_scratchpad import format_to_openai_functions
 from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser
 from langchain.chat_models import AzureChatOpenAI, ChatOpenAI
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
@@ -27,7 +26,6 @@ def get_openai_function_agent(
         [
             ("system", system_message),
             MessagesPlaceholder(variable_name="messages"),
-            MessagesPlaceholder(variable_name="agent_scratchpad"),
         ]
     )
     if tools:
@@ -36,15 +34,5 @@ def get_openai_function_agent(
         )
     else:
         llm_with_tools = llm
-    agent = (
-        {
-            "messages": lambda x: x["messages"],
-            "agent_scratchpad": lambda x: format_to_openai_functions(
-                x["intermediate_steps"]
-            ),
-        }
-        | prompt
-        | llm_with_tools
-        | OpenAIFunctionsAgentOutputParser()
-    )
+    agent = prompt | llm_with_tools | OpenAIFunctionsAgentOutputParser()
     return agent

backend/packages/gizmo-agent/gizmo_agent/agent_types/xml/agent.py

@@ -1,7 +1,6 @@
 import os
 
 import boto3
-from langchain.agents.format_scratchpad import format_xml
 from langchain.chat_models import BedrockChat, ChatAnthropic
 from langchain.schema.messages import AIMessage, HumanMessage
 from langchain.tools.render import render_text_description
@@ -61,10 +60,7 @@ def get_xml_agent(tools, system_message, bedrock=False):
     llm_with_stop = model.bind(stop=["</tool_input>"])
 
     agent = (
-        {
-            "messages": lambda x: construct_chat_history(x["messages"]),
-            "agent_scratchpad": lambda x: format_xml(x["intermediate_steps"]),
-        }
+        {"messages": lambda x: construct_chat_history(x["messages"])}
         | prompt
         | llm_with_stop
         | parse_output

backend/packages/gizmo-agent/gizmo_agent/agent_types/xml/prompts.py

@@ -33,7 +33,6 @@
     [
         ("system", template),
         MessagesPlaceholder(variable_name="messages"),
-        ("ai", "{agent_scratchpad}"),
     ]
 )
 

backend/packages/gizmo-agent/gizmo_agent/main.py

@@ -1,10 +1,7 @@
-import os
-from functools import partial
 from typing import Any, Mapping, Optional, Sequence
+from agent_executor.checkpoint import RedisCheckpoint
 
-from agent_executor import AgentExecutor
-from agent_executor.history import RunnableWithMessageHistory
-from langchain.memory import RedisChatMessageHistory
+from agent_executor.permchain import get_agent_executor
 from langchain.pydantic_v1 import BaseModel, Field
 from langchain.schema.messages import AnyMessage
 from langchain.schema.runnable import (
@@ -64,12 +61,9 @@ def __init__(
             _agent = get_xml_agent(_tools, system_message, bedrock=True)
         else:
             raise ValueError("Unexpected agent type")
-        agent_executor = AgentExecutor(
-            agent=_agent,
-            tools=_tools,
-            handle_parsing_errors=True,
-            max_iterations=10,
-        )
+        agent_executor = get_agent_executor(
+            tools=_tools, agent=_agent, checkpoint=RedisCheckpoint()
+        ).with_config({"recursion_limit": 10})
         super().__init__(
             tools=tools,
             agent=agent,
@@ -81,40 +75,33 @@ def __init__(
 
 
 class AgentInput(BaseModel):
-    input: AnyMessage
+    messages: AnyMessage
 
 
 class AgentOutput(BaseModel):
     messages: Sequence[AnyMessage] = Field(..., extra={"widget": {"type": "chat"}})
-    output: str
-
-
-agent = ConfigurableAgent(
-    agent=GizmoAgentType.GPT_35_TURBO,
-    tools=[],
-    system_message=DEFAULT_SYSTEM_MESSAGE,
-    assistant_id=None,
-).configurable_fields(
-    agent=ConfigurableField(id="agent_type", name="Agent Type"),
-    system_message=ConfigurableField(id="system_message", name="System Message"),
-    assistant_id=ConfigurableField(id="assistant_id", name="Assistant ID"),
-    tools=ConfigurableFieldMultiOption(
-        id="tools",
-        name="Tools",
-        options=TOOL_OPTIONS,
-        default=[],
-    ),
+
+
+agent = (
+    ConfigurableAgent(
+        agent=GizmoAgentType.GPT_35_TURBO,
+        tools=[],
+        system_message=DEFAULT_SYSTEM_MESSAGE,
+        assistant_id=None,
+    )
+    .configurable_fields(
+        agent=ConfigurableField(id="agent_type", name="Agent Type"),
+        system_message=ConfigurableField(id="system_message", name="System Message"),
+        assistant_id=ConfigurableField(id="assistant_id", name="Assistant ID"),
+        tools=ConfigurableFieldMultiOption(
+            id="tools",
+            name="Tools",
+            options=TOOL_OPTIONS,
+            default=[],
+        ),
+    )
+    .with_types(input_type=AgentInput, output_type=AgentOutput)
 )
-agent = RunnableWithMessageHistory(
-    agent,
-    # first arg should be a function that
-    # - accepts a single arg "session_id"
-    # - returns a BaseChatMessageHistory instance
-    partial(RedisChatMessageHistory, url=os.environ["REDIS_URL"]),
-    input_key="input",
-    output_key="messages",
-    history_key="messages",
-).with_types(input_type=AgentInput, output_type=AgentOutput)
 
 if __name__ == "__main__":
     import asyncio
@@ -123,8 +110,8 @@ class AgentOutput(BaseModel):
 
     async def run():
         async for m in agent.astream_log(
-            {"input": HumanMessage(content="whats your name")},
-            config={"configurable": {"thread_id": "test1"}},
+            {"messages": HumanMessage(content="whats your name")},
+            config={"configurable": {"user_id": "1", "thread_id": "test1"}},
         ):
             print(m)