Best practices for adding branching logic?

[codeananda]

I'm giving an LLM 37 categories and asking it to label pieces of text with the categories that apply to it (likely multiple for each text). I ask it to output its response as a markdown table.

Problem: the LLM doesn't always return answers for all the categories.

I want to check if all the categories have been returned, if they have, finish. If they haven't, ask it to classify the categories it forgot about (or, if that's too complicated, ask it to do it again).

I've thought about a RouterChain but am not sure how to handle the default chain (is there a DoNothingChain?). SequentialChain also confuses me since you cannot account for different actions based on Yes/No answer from 'does this contain all the classes?' Subclass Chain and write it?

In general though, I'm not sure how to add this branching logic?

from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate

chat = ChatOpenAI(temperature=0, model_name='gpt-4')

survey_response = 'example response'
classes = ['class 1', 'class 2', 'class 3', ...]

template_string = """Here is a response from a survey question. I am performing
thematic analysis on it and wish to classify it into one of a list of pre-defined
classes.

Response:
####{response}####

Please output whether the response falls under any of these categories.

Classes:
####{classes}####

Output should be formatted as a table with 4 columns: 1) class, 2) is_member
 1/0 depending on if the response is a member of the class), 3) confidence_score 
 (a confidence interval for whether the response does fit into that class, 
 use low (0-20%), medium (50%-80%) and high (80%+), 4) exerpt (excerpt from the 
 response that supports the classification).
"""

prompt_template = ChatPromptTemplate.from_template(template_string)

input_message = prompt_template.format_messages(
    response=survey_response,
    classes=classes
)
llm_response = chat(input_message)

[brandonwilde]

I imagine you'll first want to add some kind of output parser in order to make sure the output is being interpreted correctly. After that, you could use a RunnableBranch (https://python.langchain.com/docs/expression_language/how_to/routing) to select what should happen next. Perhaps something like this:

from langchain_core.runnables import RunnableBranch

branch = RunnableBranch(
    (lambda x: any([class in x["parsed_llm_response"] for class in classes], finish_chain),
    retry_chain,
)

full_chain = prompt_template | llm_response | output_parser | branch
full_chain.invoke({"response": "...", "classes": "..."})
``