sketching out simple chain for having agents provide feedback to each other in a langchain chain

andrew_langchain_self_critique.py

import streamlit as st
from langchain.llms import TextGen
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
import time
from langchain.chains import LLMChain, SimpleSequentialChain
from langchain import PromptTemplate


from langchain.globals import set_debug
# ANDREW -- This setting increases the log output of langchain for debug puposes. Flip to False to reduce the log output.
set_debug(True)

# Placeholder WebSocket URLs for the models
model_1_url = "ws://192.168.0.44:5005"
model_2_url = "ws://192.168.0.54:5005"

st.title("Two-Model, Three-Step LLM Interaction")

# Create Model

def create_model(model_url, temperature=0.7, max_tokens=8192):
    """
    Initializes & returns a TextGen model. ANDREW, this function will need to be adjusted
    if you want different models to be used for each actor.

    :param model_url:
    :param temperature:
    :param max_tokens:
    :return:
    """
    try:
        llm = TextGen(
            model_url=model_url,
            streaming=True,
            callbacks=[StreamingStdOutCallbackHandler()],
            temperature=temperature,
            max_tokens=max_tokens
        )
        return llm
    except Exception as e:
        # Display the error for diagnostic purposes
        st.error(f"An error occurred initializing the LLM: {e}")

with st.form("llm_interaction"):
    user_input = st.text_area("Enter your question:")
    temperature = st.slider("Select the temperature:", 0.0, 1.0, 0.7)
    max_tokens = st.slider("Select the max tokens:", 50, 8192, 250)
    submitted = st.form_submit_button("Submit")

    # Create two agents that we can use throughout, instead of ephemeral agents like before that are discarded when the function returns.
    agent_one = create_model(model_1_url, temperature, max_tokens)
    agent_two = create_model(model_2_url, temperature, max_tokens)

    if submitted and user_input:
        # Set up 1st chain for passing user input to agent 1.
        initial_template = "User Input: {original_user_input}\n In your response include user_input: <original_input>, response: <your_response>"
        first_prompt = PromptTemplate(template=initial_template, input_variables=["original_user_input"])
        first_chain = LLMChain(llm=agent_one, prompt_template=first_prompt)

        # 2nd chain for improving the answer
        improve_answer_template = "How would you improve this answer?\n {first_chain_output}  \nFormat your response as follows: user_input: <original_input>, response: <existing_response>, improvement_suggestions: <your_improvement_suggestions>"
        second_prompt = PromptTemplate(template=improve_answer_template, input_variables=["first_chain_output"])
        second_chain = LLMChain(llm=agent_two, prompt_template=second_prompt)

        # 3rd chain for refining the answer... Second chain output _should_ contain original user input, existing response, and improvement suggestions... double check in the debug logs
        refine_answer_template = "{second_chain_output}\nRefine the final answer based on the suggestions above. Your response should only include the final answer implementing the suggested edits."
        third_prompt = PromptTemplate(template=refine_answer_template, input_variables=["second_chain_output"])
        # 3rd chain re-uses the original agent. If we initialized it with a history, we could skip passing the output from the previous steps..
        third_chain = LLMChain(llm=agent_one, prompt_template=third_prompt)

        combined_chain = SimpleSequentialChain(chains=[first_chain, second_chain, third_chain])

        # Pass user input to the combined chain.
        #  - The SimpleSequentialChain will pass the output of each step to the next step. This is why we're asking the agent to rewrite the output of the first step to include the original user input.
        #  - If we initialize the 1st agent with a chat history, we could probably skip that step.
        st.write("Final refined answer:", combined_chain.run(user_input))