[Langgraph] batch processing

[Yen444]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

from langgraph.graph import END, StateGraph

class SubQuery:
    def __init__(self, retriever, llm):
        self.retriever = retriever
        self.llm = llm

    def create_retriever_chain(self):
        retriever_chain = ()  # Define the actual retriever chain logic
        return retriever_chain

    def create_llm_chain(self):
        llm_chain = ()  # Define the actual LLM chain logic
        return llm_chain

    def generate_sub_query(self, state: 'SystemState') -> 'SystemState':
        # Implement sub-query generation logic here
        return state

    def re_generate_sub_query(self, state: 'SystemState') -> 'SystemState':
        # Implement logic for re-generating sub-queries here
        return state

    def answer_sub_queries(self, state: 'SystemState') -> 'SystemState':
        retriever_chain = self.create_retriever_chain()
        llm_chain = self.create_llm_chain()

        sub_queries = state.get('sub_queries', [])

        # Assuming the chain's `batch` method processes multiple queries at once
        retriever_outputs = retriever_chain.batch(sub_queries)
        all_answers = llm_chain.batch(retriever_outputs)

        # Update the state with the answers
        state['answers'] = all_answers

        return state

    def join_answer_and_evaluate(self, state: 'SystemState') -> 'SystemState':
        # Implement logic to combine answers and perform evaluation
        return state

    def should_stop(self, state: 'SystemState') -> bool:
        # Implement logic to determine whether the workflow should stop
        return True  # Example condition

    def workflow(self):
        workflow = StateGraph('SystemState')  # Initialize the graph

        # Define the nodes
        workflow.add_node('plan', self.generate_sub_query)
        workflow.add_node('re_plan', self.re_generate_sub_query)
        workflow.add_node('execute', self.answer_sub_queries)
        workflow.add_node('evaluate', self.join_answer_and_evaluate)

        # Define the edges
        workflow.add_edge('plan', 'execute')
        workflow.add_edge('execute', 'evaluate')
        workflow.add_edge('re_plan', 'execute')

        workflow.set_entry_point('plan')

        # Add conditional logic for looping or ending the workflow
        workflow.add_conditional_edges(
            'evaluate',
            self.should_stop,
            {
                'continue': 're_plan',
                'end': END
            }
        )

        app = workflow.compile()
        return app


# Example usage
queries = []  # Batch of input queries
sub_query = SubQuery(llm=llm, retriever=retriever)
app = sub_query.workflow()  # Call instance method
results = app.batch(queries)

Description

I am building a SubQuery system using Langgraph. The system takes a user's query, generates multiple sub-queries, answers the sub-queries in parallel using batch processing, and then combines all the sub-answers. Inside the class, I use the batch method for both a retriever_chain and an llm_chain to process the sub-queries in parallel.

However, when I try to process a batch of queries outside the class using the batch method of the compiled workflow (e.g., app.batch(queries)), it seems that the batch processing inside the class and outside the class are not working together correctly. The results are not as expected. How can I resolve this issue so that batch processing works seamlessly both inside and outside the class?

System Info

langchain==0.1.20
langchain-community==0.0.38
langchain-core==0.1.52
langchain-mistralai==0.1.7
langchain-openai==0.1.5
langchain-text-splitters==0.0.2
langchainhub==0.1.15

[dosubot[bot]]

To resolve the issue where batch processing inside the SubQuery class and outside the class using the compiled workflow's batch method (e.g., app.batch(queries)) are not working together correctly, you need to ensure that the config parameter is handled correctly in both contexts.

Here is how the batch and abatch methods are implemented in the BaseChatModel class, focusing on the config parameter:

def batch(
        self,
        inputs: List[LanguageModelInput],
        config: Optional[Union[RunnableConfig, List[RunnableConfig]]] = None,
        *,
        return_exceptions: bool = False,
        **kwargs: Optional[Any],
    ) -> List[Any]:
        config = config or None
        # If <= 1 config use the underlying models batch implementation.
        if config is None or isinstance(config, dict) or len(config) <= 1:
            if isinstance(config, list):
                config = config[0]
            return self._model(config).batch(
                inputs, config=config, return_exceptions=return_exceptions, **kwargs
            )
        # If multiple configs default to Runnable.batch which uses executor to invoke
        # in parallel.
        else:
            return super().batch(
                inputs, config=config, return_exceptions=return_exceptions, **kwargs
            )

async def abatch(
        self,
        inputs: List[LanguageModelInput],
        config: Optional[Union[RunnableConfig, List[RunnableConfig]]] = None,
        *,
        return_exceptions: bool = False,
        **kwargs: Optional[Any],
    ) -> List[Any]:
        config = config or None
        # If <= 1 config use the underlying models batch implementation.
        if config is None or isinstance(config, dict) or len(config) <= 1:
            if isinstance(config, list):
                config = config[0]
            return await self._model(config).abatch(
                inputs, config=config, return_exceptions=return_exceptions, **kwargs
            )
        # If multiple configs default to Runnable.batch which uses executor to invoke
        # in parallel.
        else:
            return await super().abatch(
                inputs, config=config, return_exceptions=return_exceptions, **kwargs
            )

Key Points to Check:

1. Single vs. Multiple Configurations:

   • Ensure that when you call app.batch(queries), the config parameter is either None, a single dictionary, or a list with one item if you want to use the underlying model's batch implementation.
   • If you have multiple configurations, the method defaults to using Runnable.batch, which uses an executor to invoke in parallel.

2. Consistent Configuration Handling:

   • Make sure that the config parameter passed to app.batch(queries) is consistent with what the SubQuery class expects. If the SubQuery class expects a single configuration, ensure that you are not passing a list with multiple configurations.

3. Error Handling:

   • Use the return_exceptions parameter to handle any exceptions that might occur during batch processing. This can help in debugging issues related to batch processing.

Additionally, the LangChain framework provides helper methods like _batch_with_config and _abatch_with_config to handle configurations and callbacks effectively. These methods ensure that configurations are correctly patched and callbacks are managed properly:

def _batch_with_config(
        self,
        func: Union[
            Callable[[List[Input]], List[Union[Exception, Output]]],
            Callable[
                [List[Input], List[CallbackManagerForChainRun]],
                List[Union[Exception, Output]],
            ],
            Callable[
                [List[Input], List[CallbackManagerForChainRun], List[RunnableConfig]],
                List[Union[Exception, Output]],
            ],
        ],
        input: List[Input],
        config: Optional[Union[RunnableConfig, List[RunnableConfig]]] = None,
        *,
        return_exceptions: bool = False,
        run_type: Optional[str] = None,
        **kwargs: Optional[Any],
    ) -> List[Output]:
        """Helper method to transform an Input value to an Output value,
        with callbacks. Use this method to implement invoke() in subclasses."""
        if not input:
            return []
        configs = get_config_list(config, len(input))
        callback_managers = [get_callback_manager_for_config(c) for c in configs]
        run_managers = [
            callback_manager.on_chain_start(
                dumpd(self),
                input,
                run_type=run_type,
                name=config.get("run_name") or self.get_name(),
                run_id=config.pop("run_id", None),
            )
            for callback_manager, input, config in zip(
                callback_managers, input, configs
            )
        ]
        try:
            if accepts_config(func):
                kwargs["config"] = [
                    patch_config(c, callbacks=rm.get_child())
                    for c, rm in zip(configs, run_managers)
                ]
            if accepts_run_manager(func):
                kwargs["run_manager"] = run_managers
            output = func(input, **kwargs)  # type: ignore[call-arg]
        except BaseException as e:
            for run_manager in run_managers:
                run_manager.on_chain_error(e)
            if return_exceptions:
                return cast(List[Output], [e for _ in input])
            else:
                raise
        else:
            first_exception: Optional[Exception] = None
            for run_manager, out in zip(run_managers, output):
                if isinstance(out, Exception):
                    first_exception = first_exception or out
                    run_manager.on_chain_error(out)
                else:
                    run_manager.on_chain_end(out)
            if return_exceptions or first_exception is None:
                return cast(List[Output], output)
            else:
                raise first_exception

By ensuring that the config parameter is handled correctly and consistently in both the SubQuery class and the app.batch(queries) method, and by leveraging the helper methods provided by the LangChain framework, you should be able to resolve the issue with batch processing in your Langgraph-based system ^[1][2].

To continue talking to Dosu, mention @dosu.