diff --git a/examples/runnable-pubsub.py b/examples/runnable-pubsub.py index f3db1e18d..6331bbd84 100644 --- a/examples/runnable-pubsub.py +++ b/examples/runnable-pubsub.py @@ -1,4 +1,5 @@ from operator import itemgetter +from pprint import pprint from langchain.chat_models.openai import ChatOpenAI from langchain.prompts import SystemMessagePromptTemplate @@ -103,10 +104,16 @@ web_researcher = PubSub( connection=InMemoryPubSubConnection(), ) -for output in web_researcher.stream({"question": "What food do turtles eat?"}): - print("got output", output) +# for output in web_researcher.stream({"question": "What food do turtles eat?"}): +# print("got output", output) -print("---done---") +# print("---done with stream()---") + +pprint( + web_researcher.batch( + [{"question": "What food do turtles eat?"}, {"question": "What is art?"}] + ) +) # agent = PubSub( diff --git a/permchain/pubsub.py b/permchain/pubsub.py index ed845c459..ff170a272 100644 --- a/permchain/pubsub.py +++ b/permchain/pubsub.py @@ -1,5 +1,4 @@ from __future__ import annotations -from ctypes import Union import queue from abc import ABC @@ -16,6 +15,7 @@ from typing import ( Set, Tuple, TypeVar, + Union, ) from langchain.load.dump import dumpd @@ -78,7 +78,7 @@ class PubSub(Serializable, Runnable[Any, Any], ABC): input: Sequence[Any], config: Optional[Union[RunnableConfig, List[RunnableConfig]]] = None, ) -> Any: - configs = self._get_config_list(config) + configs = self._get_config_list(config, len(input)) with get_executor_for_config(configs[0]) as executor: return super().batch( input,