WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-07 10:17:50 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #24 from langchain-ai/harrison/langgraph

Browse Source 
start langgraph
This commit is contained in:
Nuno Campos
2024-01-07 09:56:33 -08:00
committed by GitHub
parent 80147c02e9 6875401cc9
commit 4bdca9d253
13 changed files with 2001 additions and 915 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Makefile
+1 -1
View File
@@ -18,7 +18,7 @@ test:
poetry run pytest
test_watch:
poetry run ptw
poetry run ptw --snapshot-update --now . -- -vv -x tests
######################
# LINTING AND FORMATTING
examples/.gitignore
+1
View File
@@ -0,0 +1 @@
*.db
examples/langgraph.ipynb
+399
View File
@@ -0,0 +1,399 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "396e20d9-8684-40ea-a46a-e3dfa36ed5a6",
"metadata": {},
"source": [
"## Existing Agent Executor"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "d642e6af-217a-4414-a78c-509b44155eca",
"metadata": {},
"outputs": [],
"source": [
"from langchain.chat_models import ChatOpenAI\n",
"from langchain import hub\n",
"from langchain.agents import create_openai_functions_agent\n",
"from langchain_community.chat_models import ChatOpenAI\n",
"from langchain_community.tools.tavily_search import TavilySearchResults\n",
"from langchain_core.runnables import RunnablePassthrough\n",
"from permchain.langgraph import Graph, END\n",
"\n",
"tools = [TavilySearchResults(max_results=1)]\n",
"\n",
"# Get the prompt to use - you can modify this!\n",
"prompt = hub.pull(\"hwchase17/openai-functions-agent\")\n",
"\n",
"# Choose the LLM that will drive the agent\n",
"llm = ChatOpenAI(model=\"gpt-3.5-turbo-1106\")\n",
"\n",
"# Construct the OpenAI Functions agent\n",
"agent_runnable = create_openai_functions_agent(llm, tools, prompt)\n",
"\n",
"from langchain_core.agents import AgentFinish\n",
"# Define decision-making logic\n",
"def should_continue(data):\n",
" # Logic to decide whether to continue in the loop or exit\n",
" if isinstance(data['agent_outcome'], AgentFinish):\n",
" return \"exit\"\n",
" else:\n",
" return \"continue\"\n",
" \n",
"def execute_tools(data):\n",
" agent_action = data.pop('agent_outcome')\n",
" observation = {t.name: t for t in tools}[agent_action.tool].invoke(agent_action.tool_input)\n",
" data['intermediate_steps'].append((agent_action, observation))\n",
" return data\n",
" \n",
" \n",
"\n",
"# Define agents\n",
"agent = RunnablePassthrough.assign(\n",
" agent_outcome = agent_runnable\n",
")\n",
"\n",
"\n",
"# Define a new graph\n",
"workflow = Graph()\n",
"\n",
"workflow.add_node(\"agent\", agent)\n",
"workflow.add_node(\"tools\", execute_tools)\n",
"\n",
"workflow.set_entry_point(\"agent\")\n",
"\n",
"workflow.add_conditional_edges(\n",
" \"agent\",\n",
" should_continue,\n",
" {\n",
" \"continue\": \"tools\",\n",
" \"exit\": END\n",
" }\n",
")\n",
"\n",
"workflow.add_edge('tools', 'agent')\n",
"\n",
"chain = workflow.compile()"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "c46bd262-9605-4449-9391-f6b6e0fe440e",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'input': 'what is the weather in sf',\n",
" 'intermediate_steps': [(AgentActionMessageLog(tool='tavily_search_results_json', tool_input={'query': 'weather in San Francisco'}, log=\"\\nInvoking: `tavily_search_results_json` with `{'query': 'weather in San Francisco'}`\\n\\n\\n\", message_log=[AIMessage(content='', additional_kwargs={'function_call': {'name': 'tavily_search_results_json', 'arguments': '{\"query\":\"weather in San Francisco\"}'}})]),\n",
" [{'url': 'https://www.weather25.com/north-america/usa/california/san-francisco',\n",
" 'content': 'will give you an idea of weather trends in San Francisco. For example, the weather in San Francisco in January 2024. San Francisco 14 day weather The weather today in San Francisco San Francisco weather report The weather in San Francisco, United States San Francisco weather by months San Francisco weather the weather in San Francisco including humidity, wind, chance of rain and more on the San Francisco current weather01 January 02 February 03 March 04 April 05 May 06 June 07 July 08 August 09 September 10 October 11 November 12 December. ... For example, the weather in San Francisco in January 2024. These trends can be helpful when planning trips to San Francisco or preparing for the weather in advance. There are many factors to consider when looking at the ...'}])],\n",
" 'agent_outcome': AgentFinish(return_values={'output': 'For the current weather in San Francisco, you can visit the following website: [San Francisco Weather](https://www.weather25.com/north-america/usa/california/san-francisco). This will provide you with the latest weather updates including humidity, wind, chance of rain, and more.'}, log='For the current weather in San Francisco, you can visit the following website: [San Francisco Weather](https://www.weather25.com/north-america/usa/california/san-francisco). This will provide you with the latest weather updates including humidity, wind, chance of rain, and more.')}"
]
},
"execution_count": 2,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"chain.invoke({\"input\": \"what is the weather in sf\", \"intermediate_steps\": []})"
]
},
{
"cell_type": "markdown",
"id": "592c3886-71d1-4539-80dd-111e55cc3a85",
"metadata": {},
"source": [
"## Reflexion Agent"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "f6f96e81-4a20-4599-a625-8d18df6fa76d",
"metadata": {},
"outputs": [],
"source": [
"from langchain.agents import AgentExecutor, BaseMultiActionAgent, Tool\n",
"from langchain.schema import AgentAction, AgentFinish\n",
"from langchain_core.language_models.chat_models import BaseChatModel\n",
"from langchain.chains import LLMChain\n",
"\n",
"from langchain.globals import set_llm_cache\n",
"\n",
"from dotenv import load_dotenv\n",
"\n",
"from pydantic import BaseModel\n",
"\n",
"from langchain.chat_models import ChatOpenAI\n",
"from langchain.cache import SQLiteCache\n",
"\n",
"from langchain_core.output_parsers import BaseOutputParser\n",
"\n",
"from langchain.prompts.chat import ChatPromptTemplate\n",
"from langchain.callbacks import get_openai_callback\n",
"from langchain.tools.tavily_search import TavilySearchResults\n",
"from langchain.utilities.tavily_search import TavilySearchAPIWrapper\n",
"from langchain.pydantic_v1 import BaseModel\n",
"import os\n",
"\n",
"from langchain.agents import AgentType, initialize_agent, load_tools\n",
"\n",
"set_llm_cache(SQLiteCache(database_path=\".langchain.db\"))\n",
"\n",
"\n",
"llm = ChatOpenAI(\n",
" temperature=0.0,\n",
" max_tokens=2000,\n",
" max_retries=100,\n",
" model=\"gpt-4-1106-preview\",\n",
")\n",
"\n",
"search = TavilySearchAPIWrapper()\n",
"tavily_tool = TavilySearchResults(api_wrapper=search, max_results=5)\n",
"\n",
"NEXT_STEP_TEMPLATE = \"\"\"You are expert researcher trying answer a question ~250 words. You are asked to answer the following question: {question}\n",
"\n",
"The way you are going to answer the question is as follows:\n",
"\n",
"1. Revise your previous answer using the new information.\n",
" - You should use the previous critique to add important information to your answer.\n",
" _ You MUST include numerical citations in your revised answer to ensure it can be verified.\n",
" - Add a \"References\" section to the bottom of your answer (which does not count towards the word limit). In form of:\n",
" - [1] https://example.com\n",
" - [2] https://example.com\n",
" - You should use the previous critique to remove superfluous information from your answer and make SURE it is not more than 250 words.\n",
"2. Reflect and critique your answer. Specifically, you should:\n",
" - Think about what is missing from your answer.\n",
" - Think about what is superfluous in your answer.\n",
" - Think about what search query you should use next to improve your answer.\n",
" Give your answer in exactly 2 parts. The first should address what is missing from your answer. The second should address what could be removed from your answer. Your should be VERY harsh as we really want to improve the answer.\n",
"3. Give the search query you came up with to improve your answer.\n",
"\n",
"Previous steps: \n",
"\n",
"{previous_steps}\n",
"\n",
"===\n",
"\n",
"Format your answer as follows:\n",
"\n",
"Revised answer: [give your revised answer based on the previous critique and new information from the search engine then the \"References\" section]\n",
"Critique: [give your harsh critique of your revised answer in 2 parts: what is missing and what is superfluous]\n",
"Search query: [give the new search query you came up with to enter into the search engine to improve your answer. If you have more than one, make sure they are comma separated and in quotes]\n",
"\n",
"SAY NOTHING else please.\"\"\"\n",
"\n",
"INITIAL_ANSWER_TEMPLATE = \"\"\"You are expert researcher trying answer a question ~250 words. You are asked to answer the following question: {question}\n",
"\n",
"The way you are going to answer the question is as follows:\n",
"\n",
"1. Give a detailed in ~250 words.\n",
"2. Reflect and critique your answer. Specifically, you should:\n",
" - Think about what is missing from your answer.\n",
" - Think about what is superfluous in your answer.\n",
" - Think about what search query you should use next to improve your answer.\n",
" Give your answer in exactly 2 parts. The first should address what is missing from your answer. The second should address what could be removed from your answer. Your should be VERY harsh as we really want to improve the answer.\n",
"3. Give the search query you came up with to improve your answer.\n",
"\n",
"===\n",
"\n",
"Format your answer as follows:\n",
"\n",
"Answer: [give your initial answer]\n",
"Critique: [give your harsh critique of your answer in 2 parts: what is missing and what is superfluous]\n",
"Search query: [give the search query you came up with to improve your answer. If you have more than one, make sure they are comma separated and in quotes]\n",
"\n",
"SAY NOTHING else please.\"\"\"\n",
"\n",
"\n",
"class ReflexionStep(BaseModel):\n",
" \"\"\"A single step in the reflexion process.\"\"\"\n",
"\n",
" answer: str\n",
" critique: str\n",
" search_query: str\n",
"\n",
" def __str__(self):\n",
" return f\"Answer: {self.answer}\\nCritique: {self.critique}\\nSearch query: {self.search_query}\"\n",
"\n",
"def _parse_reflexion_step(output: str) -> tuple[str, str, str]:\n",
" # find answer using .split()\n",
" if (\"Answer:\" not in output and \"Revised answer:\" not in output) or not \"Critique:\" in output or not \"Search query:\" in output:\n",
" raise ValueError(f\"The output is not formatted correctly. Output: {output}\")\n",
" if \"Answer:\" in output:\n",
" answer = output.split(\"Answer:\")[1].split(\"Critique:\")[0].strip()\n",
" else:\n",
" answer = output.split(\"Revised answer:\")[1].split(\"Critique:\")[0].strip()\n",
" critique = output.split(\"Critique:\")[1].split(\"Search query:\")[0].strip()\n",
" search_query = output.split(\"Search query:\")[1].strip()\n",
" return answer, critique, search_query\n",
"\n",
"class ReflexionStepParser(BaseOutputParser[ReflexionStep]):\n",
" \"\"\"Parser for the reflexion step.\"\"\"\n",
"\n",
" def parse(self, output: str) -> ReflexionStep:\n",
" \"\"\"Parse the output.\"\"\"\n",
" # try to find answer or initial answer\n",
" answer, critique, search_query = _parse_reflexion_step(output)\n",
" return ReflexionStep(\n",
" answer=answer, critique=critique, search_query=search_query\n",
" )"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "7708fa95-547b-4bea-b126-3656de7d5873",
"metadata": {},
"outputs": [],
"source": [
"initial_chain = RunnablePassthrough.assign(\n",
" agent_outcome = ChatPromptTemplate.from_template(INITIAL_ANSWER_TEMPLATE) | llm | ReflexionStepParser() | (lambda x: AgentAction(\n",
" tool=\"tavily_search_results_json\",\n",
" tool_input=x.search_query,\n",
" log=str(x),\n",
" ))\n",
")\n",
"\n",
"def prep_next(inputs):\n",
" intermediate_steps = inputs[\"intermediate_steps\"]\n",
" previous_steps = list[str]()\n",
"\n",
" for i, (action, observation) in enumerate(intermediate_steps, start=1):\n",
" last_step_str = f\"\"\"Step {i}:\n",
"\n",
"{action.log}\n",
"\n",
"Search output for \"{action.tool_input}\":\n",
"\n",
"{observation}\"\"\"\n",
" previous_steps.append(last_step_str)\n",
"\n",
" previous_steps_str = \"\\n\\n\".join(previous_steps)\n",
" inputs[\"previous_steps\"] = previous_steps_str\n",
" return inputs\n",
" \n",
"next_chain = RunnablePassthrough.assign(\n",
" agent_outcome = prep_next | ChatPromptTemplate.from_template(NEXT_STEP_TEMPLATE) | llm | ReflexionStepParser() | (lambda x: AgentAction(\n",
" tool=\"tavily_search_results_json\",\n",
" tool_input=x.search_query,\n",
" log=str(x),\n",
" ))\n",
")\n",
"\n",
"def finish(inputs):\n",
" intermediate_steps = inputs[\"intermediate_steps\"]\n",
" last_action, _ = intermediate_steps[-1]\n",
" last_step_str = last_action.log\n",
" # extract answer\n",
" answer, _, _ = _parse_reflexion_step(last_step_str)\n",
"\n",
" first_action, _ = intermediate_steps[0]\n",
" first_step_str = first_action.log\n",
" # extract answer\n",
" initial_answer, _, _ = _parse_reflexion_step(first_step_str)\n",
"\n",
" return AgentFinish(\n",
" log=\"Reached max steps.\",\n",
" return_values={\"output\": answer, \"initial_answer\": initial_answer},\n",
" )\n",
"\n",
"\n",
"def execute_tools(data):\n",
" agent_action = data.pop('agent_outcome')\n",
" observation = {t.name: t for t in tools}[agent_action.tool].invoke(agent_action.tool_input)\n",
" data['intermediate_steps'].append((agent_action, observation))\n",
" return data\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "d6cdd1cd-e480-4dd7-99b4-9018eb243b4d",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"AgentFinish(return_values={'output': \"The current weather in San Francisco can be accessed through various weather reporting services, which provide real-time temperature, humidity, wind, and chances of precipitation [1]. Historically, San Francisco experiences a mild, Mediterranean climate with average temperatures ranging from the low 50s to the mid-60s Fahrenheit. The city's unique topography creates microclimates, leading to significant weather variations across different neighborhoods. San Francisco's summers are notably cooler compared to other Californian cities, largely due to the cold California Current and persistent fog, especially in June and July. Winters are mild and the wettest months span from November to March, with an annual rainfall average of approximately 23 inches. Wind is a prominent feature, with spring being particularly windy. For historical weather extremes and average wind speeds, additional specific data can be sought from climatological records.\\n\\nReferences:\\n[1] https://www.weather25.com/north-america/usa/california/san-francisco\", 'initial_answer': \"The weather in San Francisco (SF) is characterized by a mild, Mediterranean-like climate with wet winters and dry summers. The city's unique topography and coastal location result in microclimates, where weather conditions can vary significantly from one neighborhood to another. Average temperatures typically range from the low 50s to the mid-60s Fahrenheit throughout the year. Summers in San Francisco are often cooler than in other parts of California due to the cold California Current offshore and the presence of fog, particularly in June and July. The fog usually burns off by the afternoon, leading to clearer skies and slightly warmer temperatures. Winters are mild and moist, with the majority of the city's rainfall occurring between November and March. Rainfall averages around 23 inches annually. Wind is also a notable feature of San Francisco's weather, with spring being the windiest season. Despite the general patterns, it's always advisable to dress in layers due to the potential for rapid weather changes.\"}, log='Reached max steps.')"
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"workflow = Graph()\n",
"\n",
"# add actors\n",
"workflow.add_node(\"initial\", initial_chain)\n",
"workflow.add_node(\"next\", next_chain)\n",
"workflow.add_node(\"finish\", finish)\n",
"workflow.add_node(\"tools\", execute_tools)\n",
"\n",
"# Enter with initial actor, then loop through tools -> next steps until finished\n",
"workflow.set_entry_point('initial')\n",
"\n",
"workflow.add_edge('initial', 'tools')\n",
"workflow.add_conditional_edges(\n",
" 'tools',\n",
" lambda x: \"exit\" if len(x['intermediate_steps']) >= 2 else \"continue\",\n",
" {\n",
" \"continue\": 'next',\n",
" \"exit\": 'finish'\n",
" }\n",
")\n",
"workflow.add_edge('next', 'tools')\n",
"workflow.set_finish_point('finish')\n",
"\n",
"chain = workflow.compile()\n",
"\n",
"chain.invoke({\"question\": \"what is the weather in sf\", \"intermediate_steps\": []})"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "9babf196-b1fd-492d-9197-96a674f5e81d",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "58ce0d58-fb00-4dc1-a12b-8fc015474611",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.5"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
permchain/__init__.py
+2
View File
@@ -1,4 +1,5 @@
from permchain.checkpoint.base import BaseCheckpointAdapter, CheckpointAt
from permchain.langgraph import Graph
from permchain.pregel import Channel, Pregel, ReservedChannels
__all__ = [
@@ -7,4 +8,5 @@ __all__ = [
"ReservedChannels",
"BaseCheckpointAdapter",
"CheckpointAt",
"Graph",
]
permchain/langgraph/__init__.py
+129
View File
@@ -0,0 +1,129 @@
from asyncio import iscoroutinefunction
from collections import defaultdict
from typing import Any, Callable, Dict, NamedTuple
from langchain_core.runnables import Runnable
from langchain_core.runnables.base import (
RunnableLambda,
RunnableLike,
coerce_to_runnable,
)
from permchain.pregel import Channel, Pregel
class Branch(NamedTuple):
condition: Callable[..., str]
ends: dict[str, str]
def runnable(self, input: Any) -> Runnable:
result = self.condition(input)
return Channel.write_to(self.ends[result])
END = "__end__"
class Graph:
def __init__(self) -> None:
self.nodes: dict[str, Runnable] = {}
self.edges = set[tuple[str, str]]()
self.branches: defaultdict[str, list[Branch]] = defaultdict(list)
def add_node(self, key: str, action: RunnableLike) -> None:
if key in self.nodes:
raise ValueError(f"Node `{key}` already present.")
if key == END:
raise ValueError(f"Node `{key}` is reserved.")
self.nodes[key] = coerce_to_runnable(action)
def add_edge(self, start_key: str, end_key: str) -> None:
if start_key not in self.nodes:
raise ValueError(f"Need to add_node `{start_key}` first")
if end_key not in self.nodes:
raise ValueError(f"Need to add_node `{end_key}` first")
# TODO: support multiple message passing
if start_key in set(start for start, _ in self.edges):
raise ValueError(f"Already found path for {start_key}")
self.edges.add((start_key, end_key))
def add_conditional_edges(
self,
start_key: str,
condition: Callable[..., str],
conditional_edge_mapping: Dict[str, str],
) -> None:
if start_key not in self.nodes:
raise ValueError(f"Need to add_node `{start_key}` first")
if iscoroutinefunction(condition):
raise ValueError("Condition cannot be a coroutine function")
self.branches[start_key].append(Branch(condition, conditional_edge_mapping))
def set_entry_point(self, key: str) -> None:
if key not in self.nodes:
raise ValueError(f"Need to add_node `{key}` first")
self.entry_point = key
def set_finish_point(self, key: str) -> None:
if key not in self.nodes:
raise ValueError(f"Need to add_node `{key}` first")
self.finish_point = key
def compile(self) -> Pregel:
################################################
# STEP 1: VALIDATE GRAPH STRUCTURE #
################################################
all_starts = (
{start for start, _ in self.edges}
| {start for start in self.branches}
| ({self.finish_point} if hasattr(self, "finish_point") else set())
)
all_ends = (
{end for _, end in self.edges}
| {
end
for branch_list in self.branches.values()
for branch in branch_list
for end in branch.ends.values()
}
| {self.entry_point}
)
for node in self.nodes:
if node not in all_ends:
raise ValueError(f"Node `{node}` is not reachable")
if node not in all_starts:
raise ValueError(f"Node `{node}` is a dead-end")
################################################
# STEP 2: CREATE GRAPH #
################################################
outgoing_edges = defaultdict(list)
for start, end in self.edges:
outgoing_edges[start].append(end)
if hasattr(self, "finish_point"):
outgoing_edges[self.finish_point].append(END)
nodes = {
key: Channel.subscribe_to(key) | node for key, node in self.nodes.items()
}
for key, edges in outgoing_edges.items():
if edges:
nodes[key] |= Channel.write_to(*edges)
for key, branches in self.branches.items():
for branch in branches:
nodes[key] |= RunnableLambda(branch.runnable, name=f"{key}_condition")
return Pregel(
nodes=nodes,
input=self.entry_point,
output=END,
)
permchain/pregel/__init__.py
+127 -59
View File
@@ -39,6 +39,7 @@ from langchain_core.runnables.utils import (
ConfigurableFieldSpec,
get_unique_config_specs,
)
from langchain_core.tracers.log_stream import LogStreamCallbackHandler
from permchain.channels.base import (
AsyncChannelsManager,
@@ -61,7 +62,7 @@ from permchain.pregel.io import map_input, map_output
from permchain.pregel.log import logger
from permchain.pregel.read import ChannelBatch, ChannelInvoke
from permchain.pregel.reserved import ReservedChannels
from permchain.pregel.validate import validate_chains_channels
from permchain.pregel.validate import validate_graph
from permchain.pregel.write import ChannelWrite
WriteValue = Union[
@@ -81,12 +82,22 @@ def _coerce_write_value(value: WriteValue) -> Runnable[Input, Output]:
class Channel:
@overload
@classmethod
def subscribe_to(cls, channels: str, key: Optional[str] = None) -> ChannelInvoke:
def subscribe_to(
cls,
channels: str,
key: Optional[str] = None,
when: Callable[[Any], bool] | None = None,
) -> ChannelInvoke:
...
@overload
@classmethod
def subscribe_to(cls, channels: Sequence[str], key: None = None) -> ChannelInvoke:
def subscribe_to(
cls,
channels: Sequence[str],
key: None = None,
when: Callable[[Any], bool] | None = None,
) -> ChannelInvoke:
...
@classmethod
@@ -134,7 +145,7 @@ class Channel:
class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
chains: Mapping[str, ChannelInvoke | ChannelBatch]
nodes: Mapping[str, ChannelInvoke | ChannelBatch]
channels: Mapping[str, BaseChannel] = Field(default_factory=dict)
@@ -153,15 +164,15 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
@root_validator(skip_on_failure=True)
def validate_pregel(cls, values: dict[str, Any]) -> dict[str, Any]:
validate_chains_channels(
values["chains"], values["channels"], values["input"], values["output"]
validate_graph(
values["nodes"], values["channels"], values["input"], values["output"]
)
return values
@property
def config_specs(self) -> list[ConfigurableFieldSpec]:
return get_unique_config_specs(
[spec for chain in self.chains.values() for spec in chain.config_specs]
[spec for node in self.nodes.values() for spec in node.config_specs]
+ (self.saver.config_specs if self.saver is not None else [])
)
@@ -205,11 +216,15 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
input: Iterator[dict[str, Any] | Any],
run_manager: CallbackManagerForChainRun,
config: RunnableConfig,
*,
output: str | Sequence[str] | None = None,
) -> Iterator[tuple[dict[str, Any] | Any, CheckpointView]]:
if config["recursion_limit"] < 1:
raise ValueError("recursion_limit must be at least 1")
# copy chains to ignore mutations during execution
processes = {**self.chains}
# assign defaults
output = output if output is not None else self.output
# copy nodes to ignore mutations during execution
processes = {**self.nodes}
# get checkpoint from saver, or create an empty one
checkpoint = self.saver.get(config) if self.saver else None
checkpoint = checkpoint or empty_checkpoint()
@@ -246,24 +261,31 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
# collect all writes to channels, without applying them yet
pending_writes = deque[tuple[str, Any]]()
# prepare tasks with config
tasks_w_config = [
(
proc,
input,
patch_config(
config,
run_name=name,
callbacks=run_manager.get_child(f"graph:step:{step}"),
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: pending_writes.extend,
CONFIG_KEY_READ: read,
},
),
)
for proc, input, name in next_tasks
]
# execute tasks, and wait for one to fail or all to finish.
# each task is independent from all other concurrent tasks
done, inflight = concurrent.futures.wait(
[
executor.submit(
proc.invoke,
input,
patch_config(
config,
callbacks=run_manager.get_child(f"pregel:step:{step}"),
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: pending_writes.extend,
CONFIG_KEY_READ: read,
},
),
)
for proc, input, _ in next_tasks
executor.submit(proc.invoke, input, config)
for proc, input, config in tasks_w_config
],
return_when=concurrent.futures.FIRST_EXCEPTION,
timeout=self.step_timeout,
@@ -283,7 +305,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
values=_updateable_channel_values(channels),
step=step + 1,
)
yield map_output(self.output, pending_writes, channels), view
yield map_output(output, pending_writes, channels), view
# if view was updated, apply writes to channels
_apply_writes_from_view(checkpoint, channels, view)
@@ -302,11 +324,24 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
input: AsyncIterator[dict[str, Any] | Any],
run_manager: AsyncCallbackManagerForChainRun,
config: RunnableConfig,
*,
output: str | Sequence[str] | None = None,
) -> AsyncIterator[tuple[dict[str, Any] | Any, CheckpointView]]:
if config["recursion_limit"] < 1:
raise ValueError("recursion_limit must be at least 1")
# copy chains to ignore mutations during execution
processes = {**self.chains}
# if running from astream_log() run each proc with streaming
do_stream = next(
(
h
for h in run_manager.handlers
if isinstance(h, LogStreamCallbackHandler)
),
None,
)
# assign defaults
output = output if output is not None else self.output
# copy nodes to ignore mutations during execution
processes = {**self.nodes}
# get checkpoint from saver, or create an empty one
checkpoint = await self.saver.aget(config) if self.saver else None
checkpoint = checkpoint or empty_checkpoint()
@@ -341,27 +376,36 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
# collect all writes to channels, without applying them yet
pending_writes = deque[tuple[str, Any]]()
# prepare tasks with config
tasks_w_config = [
(
proc,
input,
patch_config(
config,
run_name=name,
callbacks=run_manager.get_child(f"graph:step:{step}"),
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: pending_writes.extend,
CONFIG_KEY_READ: read,
},
),
)
for proc, input, name in next_tasks
]
# execute tasks, and wait for one to fail or all to finish.
# each task is independent from all other concurrent tasks
done, inflight = await asyncio.wait(
[
asyncio.create_task(
proc.ainvoke(
input,
patch_config(
config,
callbacks=run_manager.get_child(
f"pregel:step:{step}"
),
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: pending_writes.extend,
CONFIG_KEY_READ: read,
},
),
)
)
for proc, input, _ in next_tasks
asyncio.create_task(_aconsume(proc.astream(input, config)))
for proc, input, config in tasks_w_config
]
if do_stream
else [
asyncio.create_task(proc.ainvoke(input, config))
for proc, input, config in tasks_w_config
],
return_when=asyncio.FIRST_EXCEPTION,
timeout=self.step_timeout,
@@ -381,7 +425,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
values=_updateable_channel_values(channels),
step=step + 1,
)
yield map_output(self.output, pending_writes, channels), view
yield map_output(output, pending_writes, channels), view
# if view was updated, apply writes to channels
_apply_writes_from_view(checkpoint, channels, view)
@@ -399,10 +443,12 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
output: str | Sequence[str] | None = None,
**kwargs: Any,
) -> dict[str, Any] | Any:
latest: dict[str, Any] | Any = None
for chunk in self.stream(input, config, **kwargs):
for chunk in self.stream(input, config, output=output, **kwargs):
latest = chunk
return latest
@@ -410,30 +456,36 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
output: str | Sequence[str] | None = None,
**kwargs: Any,
) -> Iterator[dict[str, Any] | Any]:
return self.transform(iter([input]), config, **kwargs)
return self.transform(iter([input]), config, output=output, **kwargs)
def transform(
self,
input: Iterator[dict[str, Any] | Any],
config: RunnableConfig | None = None,
*,
output: str | Sequence[str] | None = None,
**kwargs: Any | None,
) -> Iterator[dict[str, Any] | Any]:
for output, _ in self._transform_stream_with_config(
input, self._transform, config, **kwargs
for out, _ in self._transform_stream_with_config(
input, self._transform, config, output=output, **kwargs
):
if output is not None:
yield output
if out is not None:
yield cast(dict[str, Any] | Any, out)
def step(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
output: str | Sequence[str] | None = None,
**kwargs: Any,
) -> Iterator[tuple[dict[str, Any] | Any, CheckpointView]]:
for tup in self._transform_stream_with_config(
iter([input]), self._transform, config, **kwargs
iter([input]), self._transform, config, output=output, **kwargs
):
yield cast(tuple[dict[str, Any] | Any, CheckpointView], tup)
@@ -441,10 +493,12 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
output: str | Sequence[str] | None = None,
**kwargs: Any,
) -> dict[str, Any] | Any:
latest: dict[str, Any] | Any = None
async for chunk in self.astream(input, config, **kwargs):
async for chunk in self.astream(input, config, output=output, **kwargs):
latest = chunk
return latest
@@ -452,37 +506,45 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
output: str | Sequence[str] | None = None,
**kwargs: Any,
) -> AsyncIterator[dict[str, Any] | Any]:
async def input_stream() -> AsyncIterator[dict[str, Any] | Any]:
yield input
async for chunk in self.atransform(input_stream(), config, **kwargs):
async for chunk in self.atransform(
input_stream(), config, output=output, **kwargs
):
yield chunk
async def atransform(
self,
input: AsyncIterator[dict[str, Any] | Any],
config: RunnableConfig | None = None,
*,
output: str | Sequence[str] | None = None,
**kwargs: Any | None,
) -> AsyncIterator[dict[str, Any] | Any]:
async for output, _ in self._atransform_stream_with_config(
input, self._atransform, config, **kwargs
async for out, _ in self._atransform_stream_with_config(
input, self._atransform, config, output=output, **kwargs
):
if output is not None:
yield output
if out is not None:
yield out
async def astep(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
output: str | Sequence[str] | None = None,
**kwargs: Any,
) -> AsyncIterator[tuple[dict[str, Any] | Any, CheckpointView]]:
async def input_stream() -> AsyncIterator[dict[str, Any] | Any]:
yield input
async for tup in self._atransform_stream_with_config(
input_stream(), self._atransform, config, **kwargs
input_stream(), self._atransform, config, output=output, **kwargs
):
yield cast(tuple[dict[str, Any] | Any, CheckpointView], tup)
@@ -633,9 +695,15 @@ def _updateable_channel_values(channels: Mapping[str, BaseChannel]) -> dict[str,
"""Return a dictionary of updateable channel values."""
values: dict[str, Any] = {}
for k, v in channels.items():
if isinstance(v, LastValue):
if isinstance(v, LastValue) and k not in [c.value for c in ReservedChannels]:
try:
values[k] = v.get()
except EmptyChannelError:
pass
return values
async def _aconsume(iterator: AsyncIterator[Any]) -> None:
"""Consume an async iterator."""
async for _ in iterator:
pass
permchain/pregel/read.py
+1
View File
@@ -40,6 +40,7 @@ class ChannelRead(RunnableLambda):
def __init__(self, channel: str) -> None:
super().__init__(func=self._read, afunc=self._aread)
self.channel = channel
self.name = f"ChannelRead<{channel}>"
def _read(self, _: Any, config: RunnableConfig) -> Any:
try:
permchain/pregel/validate.py
+10 -10
View File
@@ -6,21 +6,21 @@ from permchain.pregel.read import ChannelBatch, ChannelInvoke
from permchain.pregel.reserved import ReservedChannels
def validate_chains_channels(
chains: Mapping[str, ChannelInvoke | ChannelBatch],
def validate_graph(
nodes: Mapping[str, ChannelInvoke | ChannelBatch],
channels: dict[str, BaseChannel],
input: str | Sequence[str],
output: str | Sequence[str],
) -> None:
subscribed_channels = set[str]()
for chain in chains.values():
if isinstance(chain, ChannelInvoke):
subscribed_channels.update(chain.channels.values())
elif isinstance(chain, ChannelBatch):
subscribed_channels.add(chain.channel)
for node in nodes.values():
if isinstance(node, ChannelInvoke):
subscribed_channels.update(node.channels.values())
elif isinstance(node, ChannelBatch):
subscribed_channels.add(node.channel)
else:
raise TypeError(
f"Invalid chain type {type(chain)}, expected Channel.subscribe_to() or Channel.subscribe_to_each()"
f"Invalid node type {type(node)}, expected Channel.subscribe_to() or Channel.subscribe_to_each()"
)
for chan in subscribed_channels:
@@ -31,14 +31,14 @@ def validate_chains_channels(
if input not in channels:
channels[input] = LastValue(Any) # type: ignore[arg-type]
if input not in subscribed_channels:
raise ValueError(f"Input channel {input} is not subscribed to by any chain")
raise ValueError(f"Input channel {input} is not subscribed to by any node")
else:
for chan in input:
if chan not in channels:
channels[chan] = LastValue(Any) # type: ignore[arg-type]
if all(chan not in subscribed_channels for chan in input):
raise ValueError(
f"None of the input channels {input} are subscribed to by any chain"
f"None of the input channels {input} are subscribed to by any node"
)
if isinstance(output, str):
permchain/pregel/write.py
+1
View File
@@ -30,6 +30,7 @@ class ChannelWrite(RunnablePassthrough):
channels: Sequence[tuple[str, Runnable | None]],
):
super().__init__(func=self._write, afunc=self._awrite, channels=channels)
self.name = f"ChannelWrite<{','.join(chan for chan, _ in self.channels)}>"
@property
def config_specs(self) -> list[ConfigurableFieldSpec]:
poetry.lock
Generated
+722 -703
View File
File diff suppressed because it is too large Load Diff
pyproject.toml
+4 -3
View File
@@ -22,8 +22,8 @@ pytest-dotenv = "^0.5.2"
pytest-asyncio = "^0.20.3"
pytest-mock = "^3.10.0"
syrupy = "^4.0.2"
pytest-watch = "^4.2.0"
httpx = "^0.26.0"
pytest-watcher = "^0.3.4"
[tool.poetry.group.lint.dependencies]
ruff = "^0.1.4"
@@ -37,7 +37,8 @@ optional = true
[tool.poetry.group.dev.dependencies]
jupyter = "^1.0.0"
openai = "^0.27.8"
langchain = "^0.0.352"
langchain = "^0.1.0"
langchainhub = "^0.1.14"
[tool.ruff]
select = [ "E", "F", "I" ]
@@ -66,6 +67,6 @@ asyncio_mode = "auto"
#
# https://github.com/tophat/syrupy
# --snapshot-warn-unused Prints a warning on unused snapshots rather than fail the test suite.
addopts = "-x --full-trace --strict-markers --strict-config --durations=5 --snapshot-warn-unused"
addopts = "--full-trace --strict-markers --strict-config --durations=5 --snapshot-warn-unused"
# Registering custom markers.
# https://docs.pytest.org/en/7.1.x/example/markers.html#registering-markers
tests/test_pregel.py
+300 -71
View File
@@ -8,7 +8,7 @@ import pytest
from langchain_core.runnables import RunnablePassthrough
from pytest_mock import MockerFixture
from permchain import Channel, Pregel
from permchain import Channel, Graph, Pregel
from permchain.channels.base import InvalidUpdateError
from permchain.channels.binop import BinaryOperatorAggregate
from permchain.channels.context import Context
@@ -23,7 +23,7 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={
nodes={
"one": chain,
},
channels={
@@ -33,18 +33,26 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
input="input",
output="output",
)
graph = Graph()
graph.add_node("add_one", add_one)
graph.set_entry_point("add_one")
graph.set_finish_point("add_one")
gapp = graph.compile()
assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"}
assert app.invoke(2) == 3
assert app.invoke(2, output=["output"]) == {"output": 3}
assert repr(app), "does not raise recursion error"
assert gapp.invoke(2) == 3
def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(chains={"one": chain})
app = Pregel(nodes={"one": chain})
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.output_schema.schema() == {"title": "PregelOutput"}
@@ -59,7 +67,7 @@ def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
| Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1)
)
app = Pregel(chains={"one": chain}, output=["output", "fixed", "output_plus_one"])
app = Pregel(nodes={"one": chain}, output=["output", "fixed", "output_plus_one"])
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.output_schema.schema() == {
@@ -83,7 +91,7 @@ def test_invoke_single_process_in_out_reserved_is_last(mocker: MockerFixture) ->
| Channel.write_to("output")
)
app = Pregel(chains={"one": chain})
app = Pregel(nodes={"one": chain})
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.output_schema.schema() == {"title": "PregelOutput"}
@@ -96,7 +104,7 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={
nodes={
"one": chain,
},
output=["output"],
@@ -116,7 +124,7 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={
nodes={
"one": chain,
},
input=["input"],
@@ -138,11 +146,11 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
nodes={"one": one, "two": two},
)
assert app.invoke(2) == 4
@@ -152,7 +160,6 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
assert view.values == {
"inbox": 3,
"input": 2,
"is_last_step": False,
}
assert output is None
elif view.step == 2:
@@ -160,7 +167,6 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
"output": 4,
"inbox": 3,
"input": 2,
"is_last_step": False,
}
assert output == 4
@@ -169,7 +175,6 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
assert view.values == {
"inbox": 3,
"input": 2,
"is_last_step": False,
}
assert output is None
# modify inbox value
@@ -179,7 +184,49 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
"output": 6,
"inbox": 5,
"input": 2,
"is_last_step": False,
}
# output is different now
assert output == 6
graph = Graph()
graph.add_node("add_one", add_one)
graph.add_node("add_one_more", add_one)
graph.set_entry_point("add_one")
graph.set_finish_point("add_one_more")
graph.add_edge("add_one", "add_one_more")
gapp = graph.compile()
assert gapp.invoke(2) == 4
for output, view in gapp.step(2):
if view.step == 1:
assert view.values == {
"add_one": 2,
"add_one_more": 3,
}
assert output is None
elif view.step == 2:
assert view.values == {
"add_one": 2,
"add_one_more": 3,
"__end__": 4,
}
assert output == 4
for output, view in gapp.step(2):
if view.step == 1:
assert view.values == {
"add_one": 2,
"add_one_more": 3,
}
assert output is None
# modify inbox value
view.values["add_one_more"] = 5
elif view.step == 2:
assert view.values == {
"add_one": 2,
"add_one_more": 5,
"__end__": 6,
}
# output is different now
assert output == 6
@@ -187,18 +234,20 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_two = (
Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
)
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
nodes={"one": one, "two": two},
channels={"inbox": Topic(int)},
input=["input", "inbox"],
)
assert [*app.stream({"input": 2, "inbox": 12})] == [13, 4] # [12 + 1, 2 + 1 + 1]
assert [*app.stream({"input": 2, "inbox": 12}, output=["output"])] == [
{"output": 13},
{"output": 4},
]
def test_batch_two_processes_in_out() -> None:
@@ -206,30 +255,43 @@ def test_batch_two_processes_in_out() -> None:
time.sleep(inp / 10)
return inp + 1
chain_one = (
Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one")
)
chain_two = (
Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output")
)
one = Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one")
two = Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output")
app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two})
app = Pregel(nodes={"one": one, "two": two})
assert app.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
assert app.batch([3, 2, 1, 3, 5], output=["output"]) == [
{"output": 5},
{"output": 4},
{"output": 3},
{"output": 5},
{"output": 7},
]
graph = Graph()
graph.add_node("add_one", add_one_with_delay)
graph.add_node("add_one_more", add_one_with_delay)
graph.set_entry_point("add_one")
graph.set_finish_point("add_one_more")
graph.add_edge("add_one", "add_one_more")
gapp = graph.compile()
assert gapp.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
test_size = 100
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
for i in range(test_size - 2):
chains[str(i)] = (
nodes[str(i)] = (
Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
)
chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
app = Pregel(chains=chains)
app = Pregel(nodes=nodes)
for _ in range(10):
assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size
@@ -244,14 +306,14 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
test_size = 100
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
for i in range(test_size - 2):
chains[str(i)] = (
nodes[str(i)] = (
Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
)
chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
app = Pregel(chains=chains)
app = Pregel(nodes=nodes)
for _ in range(3):
assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
@@ -275,10 +337,10 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two})
app = Pregel(nodes={"one": one, "two": two})
with pytest.raises(InvalidUpdateError):
# LastValue channels can only be updated once per iteration
@@ -288,11 +350,11 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N
def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
nodes={"one": one, "two": two},
channels={"output": Topic(int)},
)
@@ -308,7 +370,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None:
raise ValueError("Input is too large")
return input
chain_one = (
one = (
Channel.subscribe_to(["input"]).join(["total"])
| add_one
| Channel.write_to("output", "total")
@@ -318,7 +380,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None:
memory = MemoryCheckpoint()
app = Pregel(
chains={"chain_one": chain_one},
nodes={"one": one},
channels={"total": BinaryOperatorAggregate(int, operator.add)},
saver=memory,
)
@@ -354,15 +416,15 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None
add_one = mocker.Mock(side_effect=lambda x: x + 1)
add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_four = (
Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output")
)
app = Pregel(
chains={
"chain_one": chain_one,
nodes={
"one": one,
"chain_three": chain_three,
"chain_four": chain_four,
},
@@ -384,17 +446,17 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
inner_app = Pregel(
chains={
nodes={
"one": Channel.subscribe_to("input") | add_one | Channel.write_to("output")
}
)
chain_one = (
one = (
Channel.subscribe_to("input")
| add_10_each
| Channel.write_to("inbox_one").map()
)
chain_two = (
two = (
Channel.subscribe_to("inbox_one")
| inner_app.map()
| sorted
@@ -403,9 +465,9 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output")
app = Pregel(
chains={
"chain_one": chain_one,
"chain_two": chain_two,
nodes={
"one": one,
"two": two,
"chain_three": chain_three,
},
channels={"inbox_one": Topic(int)},
@@ -421,28 +483,24 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = (
one = (
Channel.subscribe_to("input")
| add_one
| Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough())
)
chain_two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
)
app = Pregel(nodes={"one": one, "two": two})
assert [c for c in app.stream(2)] == [3, 4]
def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("between")
chain_two = Channel.subscribe_to("between") | add_one
one = Channel.subscribe_to("input") | add_one | Channel.write_to("between")
two = Channel.subscribe_to("between") | add_one
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
)
app = Pregel(nodes={"one": one, "two": two})
# It finishes executing (once no more messages being published)
# but returns nothing, as nothing was published to OUT topic
@@ -452,13 +510,11 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
chain_two = Channel.subscribe_to("between") | add_one
one = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
two = Channel.subscribe_to("between") | add_one
with pytest.raises(ValueError):
Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
)
Pregel(nodes={"one": one, "two": two})
def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
@@ -474,13 +530,11 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
cleanup()
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_two = (
Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
)
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
nodes={"one": one, "two": two},
channels={
"inbox": Topic(int),
"ctx": Context(an_int, typ=int),
@@ -500,3 +554,178 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
else:
assert False, "Expected only two chunks"
assert cleanup.call_count == 1, "Expected cleanup to be called once"
def test_conditional_graph() -> None:
from copy import deepcopy
from langchain.llms.fake import FakeStreamingListLLM
from langchain_community.tools import tool
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from permchain.langgraph import END
# Assemble the tools
@tool()
def search_api(query: str) -> str:
"""Searches the API for the query."""
return f"result for {query}"
tools = [search_api]
# Construct the agent
prompt = PromptTemplate.from_template("Hello!")
llm = FakeStreamingListLLM(
responses=[
"tool:search_api:query",
"tool:search_api:another",
"finish:answer",
]
)
def agent_parser(input: str) -> AgentFinish | AgentAction:
if input.startswith("finish"):
_, answer = input.split(":")
return AgentFinish(return_values={"answer": answer}, log=input)
else:
_, tool_name, tool_input = input.split(":")
return AgentAction(tool=tool_name, tool_input=tool_input, log=input)
agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser)
# Define tool execution logic
def execute_tools(data: dict) -> dict:
agent_action: AgentAction = data.pop("agent_outcome")
observation = {t.name: t for t in tools}[agent_action.tool].invoke(
agent_action.tool_input
)
if data.get("intermediate_steps") is None:
data["intermediate_steps"] = []
data["intermediate_steps"].append((agent_action, observation))
return data
# Define decision-making logic
def should_continue(data: dict) -> str:
# Logic to decide whether to continue in the loop or exit
if isinstance(data["agent_outcome"], AgentFinish):
return "exit"
else:
return "continue"
# Define a new graph
workflow = Graph()
workflow.add_node("agent", agent)
workflow.add_node("tools", execute_tools)
workflow.set_entry_point("agent")
workflow.add_conditional_edges(
"agent", should_continue, {"continue": "tools", "exit": END}
)
workflow.add_edge("tools", "agent")
app = workflow.compile()
assert app.invoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
),
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
),
],
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
}
assert [
deepcopy(c)
for c in app.stream(
{"input": "what is weather in sf"}, output=["agent", "tools"]
)
] == [
{
"tools": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
{
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
}
},
{
"tools": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
),
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
),
],
}
},
]
tests/test_pregel_async.py
+304 -68
View File
@@ -7,7 +7,7 @@ import pytest
from langchain_core.runnables import RunnablePassthrough
from pytest_mock import MockerFixture
from permchain import Channel, Pregel
from permchain import Channel, Graph, Pregel
from permchain.channels.base import InvalidUpdateError
from permchain.channels.binop import BinaryOperatorAggregate
from permchain.channels.context import Context
@@ -22,7 +22,7 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={
nodes={
"one": chain,
},
channels={
@@ -32,19 +32,27 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
input="input",
output="output",
)
graph = Graph()
graph.add_node("add_one", add_one)
graph.set_entry_point("add_one")
graph.set_finish_point("add_one")
gapp = graph.compile()
assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"}
assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"}
assert await app.ainvoke(2) == 3
assert await app.ainvoke(2, output=["output"]) == {"output": 3}
assert await gapp.ainvoke(2) == 3
async def test_invoke_single_process_in_out_implicit_channels(
mocker: MockerFixture
mocker: MockerFixture,
) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(chains={"one": chain})
app = Pregel(nodes={"one": chain})
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.output_schema.schema() == {"title": "PregelOutput"}
@@ -59,7 +67,7 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N
| Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1)
)
app = Pregel(chains={"one": chain}, output=["output", "fixed", "output_plus_one"])
app = Pregel(nodes={"one": chain}, output=["output", "fixed", "output_plus_one"])
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.output_schema.schema() == {
@@ -75,7 +83,7 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N
async def test_invoke_single_process_in_out_reserved_is_last(
mocker: MockerFixture
mocker: MockerFixture,
) -> None:
add_one = mocker.Mock(side_effect=lambda x: {**x, "input": x["input"] + 1})
@@ -85,7 +93,7 @@ async def test_invoke_single_process_in_out_reserved_is_last(
| Channel.write_to("output")
)
app = Pregel(chains={"one": chain})
app = Pregel(nodes={"one": chain})
assert app.input_schema.schema() == {"title": "PregelInput"}
assert app.output_schema.schema() == {"title": "PregelOutput"}
@@ -101,9 +109,7 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={
"one": chain,
},
nodes={"one": chain},
output=["output"],
)
@@ -121,7 +127,7 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) ->
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={
nodes={
"one": chain,
},
input=["input"],
@@ -143,12 +149,10 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) ->
async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
)
app = Pregel(nodes={"one": one, "two": two})
assert await app.ainvoke(2) == 4
@@ -157,7 +161,6 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
assert view.values == {
"inbox": 3,
"input": 2,
"is_last_step": False,
}
assert output is None
elif view.step == 2:
@@ -165,7 +168,6 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
"output": 4,
"inbox": 3,
"input": 2,
"is_last_step": False,
}
assert output == 4
@@ -174,7 +176,6 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
assert view.values == {
"inbox": 3,
"input": 2,
"is_last_step": False,
}
assert output is None
# modify inbox value
@@ -184,7 +185,49 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
"output": 6,
"inbox": 5,
"input": 2,
"is_last_step": False,
}
# output is different now
assert output == 6
graph = Graph()
graph.add_node("add_one", add_one)
graph.add_node("add_one_more", add_one)
graph.set_entry_point("add_one")
graph.set_finish_point("add_one_more")
graph.add_edge("add_one", "add_one_more")
gapp = graph.compile()
assert await gapp.ainvoke(2) == 4
async for output, view in gapp.astep(2):
if view.step == 1:
assert view.values == {
"add_one": 2,
"add_one_more": 3,
}
assert output is None
elif view.step == 2:
assert view.values == {
"add_one": 2,
"add_one_more": 3,
"__end__": 4,
}
assert output == 4
async for output, view in gapp.astep(2):
if view.step == 1:
assert view.values == {
"add_one": 2,
"add_one_more": 3,
}
assert output is None
# modify inbox value
view.values["add_one_more"] = 5
elif view.step == 2:
assert view.values == {
"add_one": 2,
"add_one_more": 5,
"__end__": 6,
}
# output is different now
assert output == 6
@@ -192,19 +235,20 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_two = (
Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
)
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
pubsub = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
nodes={"one": one, "two": two},
channels={"inbox": Topic(int)},
input=["input", "inbox"],
)
# [12 + 1, 2 + 1 + 1]
assert [c async for c in pubsub.astream({"input": 2, "inbox": 12})] == [13, 4]
assert [
c async for c in pubsub.astream({"input": 2, "inbox": 12}, output=["output"])
] == [{"output": 13}, {"output": 4}]
async def test_batch_two_processes_in_out() -> None:
@@ -212,33 +256,46 @@ async def test_batch_two_processes_in_out() -> None:
await asyncio.sleep(inp / 10)
return inp + 1
chain_one = (
Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one")
)
chain_two = (
Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output")
)
one = Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one")
two = Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
nodes={"one": one, "two": two},
channels={"one": LastValue(int)},
)
assert await app.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
assert await app.abatch([3, 2, 1, 3, 5], output=["output"]) == [
{"output": 5},
{"output": 4},
{"output": 3},
{"output": 5},
{"output": 7},
]
graph = Graph()
graph.add_node("add_one", add_one_with_delay)
graph.add_node("add_one_more", add_one_with_delay)
graph.set_entry_point("add_one")
graph.set_finish_point("add_one_more")
graph.add_edge("add_one", "add_one_more")
gapp = graph.compile()
assert await gapp.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
test_size = 100
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
for i in range(test_size - 2):
chains[str(i)] = (
nodes[str(i)] = (
Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
)
chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
app = Pregel(chains=chains)
app = Pregel(nodes=nodes)
# No state is left over from previous invocations
for _ in range(10):
@@ -254,14 +311,14 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
test_size = 100
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
for i in range(test_size - 2):
chains[str(i)] = (
nodes[str(i)] = (
Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
)
chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
app = Pregel(chains=chains)
app = Pregel(nodes=nodes)
# No state is left over from previous invocations
for _ in range(3):
@@ -288,10 +345,10 @@ async def test_invoke_two_processes_two_in_two_out_invalid(
) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two})
app = Pregel(nodes={"one": one, "two": two})
with pytest.raises(InvalidUpdateError):
# LastValue channels can only be updated once per iteration
@@ -301,11 +358,11 @@ async def test_invoke_two_processes_two_in_two_out_invalid(
async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
nodes={"one": one, "two": two},
channels={"output": Topic(int)},
)
@@ -321,7 +378,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
raise ValueError("Input is too large")
return input
chain_one = (
one = (
Channel.subscribe_to(["input"]).join(["total"])
| add_one
| Channel.write_to("output", "total")
@@ -331,7 +388,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
memory = MemoryCheckpoint()
app = Pregel(
chains={"chain_one": chain_one},
nodes={"one": one},
channels={"total": BinaryOperatorAggregate(int, operator.add)},
saver=memory,
)
@@ -367,15 +424,15 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -
add_one = mocker.Mock(side_effect=lambda x: x + 1)
add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_four = (
Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output")
)
app = Pregel(
chains={
"chain_one": chain_one,
nodes={
"one": one,
"chain_three": chain_three,
"chain_four": chain_four,
},
@@ -398,17 +455,17 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None
add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
inner_app = Pregel(
chains={
nodes={
"one": Channel.subscribe_to("input") | add_one | Channel.write_to("output")
}
)
chain_one = (
one = (
Channel.subscribe_to("input")
| add_10_each
| Channel.write_to("inbox_one").map()
)
chain_two = (
two = (
Channel.subscribe_to("inbox_one")
| inner_app.map()
| sorted
@@ -417,9 +474,9 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None
chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output")
app = Pregel(
chains={
"chain_one": chain_one,
"chain_two": chain_two,
nodes={
"one": one,
"two": two,
"chain_three": chain_three,
},
channels={
@@ -440,14 +497,14 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None
async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = (
one = (
Channel.subscribe_to("input")
| add_one
| Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough())
)
chain_two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two})
app = Pregel(nodes={"one": one, "two": two})
# Then invoke pubsub
assert [c async for c in app.astream(2)] == [3, 4]
@@ -455,10 +512,10 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("between")
chain_two = Channel.subscribe_to("between") | add_one
one = Channel.subscribe_to("input") | add_one | Channel.write_to("between")
two = Channel.subscribe_to("between") | add_one
app = Pregel(chains={"chain_one": chain_one, "chain_two": chain_two})
app = Pregel(nodes={"one": one, "two": two})
# It finishes executing (once no more messages being published)
# but returns nothing, as nothing was published to "output" topic
@@ -488,13 +545,11 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
cleanup_async()
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_two = (
Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
)
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
nodes={"one": one, "two": two},
channels={
"inbox": Topic(int),
"ctx": Context(an_int, an_int_async, typ=int),
@@ -527,3 +582,184 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
assert cleanup_sync.call_count == 0
assert setup_async.call_count == 1, "Expected setup to be called once"
assert cleanup_async.call_count == 1, "Expected cleanup to be called once"
async def test_conditional_graph() -> None:
from copy import deepcopy
from langchain.llms.fake import FakeStreamingListLLM
from langchain_community.tools import tool
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from permchain.langgraph import END
# Assemble the tools
@tool()
def search_api(query: str) -> str:
"""Searches the API for the query."""
return f"result for {query}"
tools = [search_api]
# Construct the agent
prompt = PromptTemplate.from_template("Hello!")
llm = FakeStreamingListLLM(
responses=[
"tool:search_api:query",
"tool:search_api:another",
"finish:answer",
]
)
async def agent_parser(input: str) -> AgentFinish | AgentAction:
if input.startswith("finish"):
_, answer = input.split(":")
return AgentFinish(return_values={"answer": answer}, log=input)
else:
_, tool_name, tool_input = input.split(":")
return AgentAction(tool=tool_name, tool_input=tool_input, log=input)
agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser)
# Define tool execution logic
async def execute_tools(data: dict) -> dict:
agent_action: AgentAction = data.pop("agent_outcome")
observation = await {t.name: t for t in tools}[agent_action.tool].ainvoke(
agent_action.tool_input
)
if data.get("intermediate_steps") is None:
data["intermediate_steps"] = []
data["intermediate_steps"].append((agent_action, observation))
return data
# Define decision-making logic
def should_continue(data: dict) -> str:
# Logic to decide whether to continue in the loop or exit
if isinstance(data["agent_outcome"], AgentFinish):
return "exit"
else:
return "continue"
# Define a new graph
workflow = Graph()
workflow.add_node("agent", agent)
workflow.add_node("tools", execute_tools)
workflow.set_entry_point("agent")
workflow.add_conditional_edges(
"agent", should_continue, {"continue": "tools", "exit": END}
)
workflow.add_edge("tools", "agent")
app = workflow.compile()
assert await app.ainvoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
),
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
),
],
"agent_outcome": AgentFinish(
return_values={"answer": "answer"}, log="finish:answer"
),
}
assert [
deepcopy(c)
async for c in app.astream(
{"input": "what is weather in sf"}, output=["agent", "tools"]
)
] == [
{
"tools": {
"input": "what is weather in sf",
"agent_outcome": AgentAction(
tool="search_api", tool_input="query", log="tool:search_api:query"
),
}
},
{
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
}
},
{
"tools": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
)
],
"agent_outcome": AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
}
},
{
"agent": {
"input": "what is weather in sf",
"intermediate_steps": [
(
AgentAction(
tool="search_api",
tool_input="query",
log="tool:search_api:query",
),
"result for query",
),
(
AgentAction(
tool="search_api",
tool_input="another",
log="tool:search_api:another",
),
"result for another",
),
],
}
},
]
patches = [c async for c in app.astream_log({"input": "what is weather in sf"})]
patch_paths = {op["path"] for log in patches for op in log.ops}
# Check that agent (one of the nodes) has its output streamed to the logs
assert "/logs/agent/streamed_output/-" in patch_paths