setting up a 2-model strategy

app.py

@@ -4,54 +4,50 @@ import networkx as nx
 import matplotlib.pyplot as plt
 from transformers import pipeline
 
-# Initialize three text2text-generation pipelines using Flan-T5-large.
-model1_generator = pipeline('text2text-generation', model='google/flan-t5-large')
-model2_generator = pipeline('text2text-generation', model='google/flan-t5-large')
-model3_generator = pipeline('text2text-generation', model='google/flan-t5-large')
+# Initialize two text2text-generation pipelines using Flan-T5-large.
+model_request_to_graph = pipeline('text2text-generation', model='google/flan-t5-large')
+model_reply_with_graph = pipeline('text2text-generation', model='google/flan-t5-large')
 
-def model1_translate_to_graph(request_text):
+def generate_input_graph(request_text):
     """
-    Stage 1: Translate the customer request into a structured graph.
-    Output valid JSON with two keys: 'nodes' and 'edges'. 
-    Nodes is a list of unique node names, and edges is a list of pairs [source, target].
+    Model 1: Convert the customer request into a structured graph.
+    The prompt instructs the model to output only valid JSON with exactly two keys:
+    - 'nodes': a list of unique node names (strings)
+    - 'edges': a list of pairs [source, target] representing directed edges.
+    The output JSON is intended to be directly used for constructing a NetworkX graph.
     """
     prompt = (
-        "Translate the following customer request into a structured graph. "
-        "Output only valid JSON with exactly two keys: 'nodes' and 'edges'. "
-        "'nodes' should be a list of unique node names, and 'edges' should be a list of pairs "
-        "of node names representing directed edges. Do not include any extra text.\n"
-        f"Customer Request: \"{request_text}\"\n"
+        "You are an assistant that converts a natural language customer request into a "
+        "structured graph. Please output only valid JSON with exactly two keys: 'nodes' and 'edges'. "
+        "'nodes' should be a list of unique strings representing node names. "
+        "'edges' should be a list of pairs [source, target] representing directed edges between nodes. "
+        "This structure is intended for use with the Python NetworkX library. "
+        "Do not include any extra text or commentary.\n\n"
+        f"Customer Request: \"{request_text}\"\n\n"
         "Structured Graph JSON:"
     )
-    output = model1_generator(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)[0]['generated_text']
+    output = model_request_to_graph(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)[0]['generated_text']
     return output
 
-def model2_generate_response(graph_description):
+def generate_reply_and_graph(request_text):
     """
-    Stage 2: Generate a detailed response based on the structured graph description.
+    Model 2: Generate a detailed response in a structured way.
+    The prompt instructs the model to produce a reply that not only answers the customer request,
+    but also provides a corresponding structured graph (as valid JSON with keys 'nodes' and 'edges')
+    that reflects the underlying relationships in the response. This graph is designed to be used 
+    with NetworkX for transparent input and output.
     """
     prompt = (
-        "Based on the following structured customer request graph (formatted as JSON), "
-        "provide a detailed and helpful response.\n"
-        f"Structured Graph: {graph_description}\n"
-        "Response:"
+        "You are an assistant that responds to a customer request in a detailed and structured manner. "
+        "First, provide a helpful textual reply to the request. Then, convert your reply into a structured graph "
+        "that shows the main concepts and relationships. Output only valid JSON after your reply, with exactly two keys: "
+        "'nodes' (a list of unique strings representing node names) and 'edges' (a list of pairs [source, target] representing "
+        "directed edges between nodes). This graph will be used with the Python NetworkX library. "
+        "Do not include any extra text or commentary outside of the JSON.\n\n"
+        f"Customer Request: \"{request_text}\"\n\n"
+        "Detailed Response and Structured Graph JSON:"
     )
-    output = model2_generator(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)[0]['generated_text']
-    return output
-
-def model3_response_to_graph(response_text):
-    """
-    Stage 3: Convert the detailed response into a structured graph.
-    Output only valid JSON with two keys: 'nodes' and 'edges'.
-    """
-    prompt = (
-        "Translate the following response into a structured graph. "
-        "Output only valid JSON with two keys: 'nodes' (a list of concept names) and 'edges' "
-        "(a list of pairs representing relationships between these concepts). Do not include any extra text.\n"
-        f"Response: \"{response_text}\"\n"
-        "Structured Response Graph JSON:"
-    )
-    output = model3_generator(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)[0]['generated_text']
+    output = model_reply_with_graph(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)[0]['generated_text']
     return output
 
 def try_parse_json(text):
@@ -88,12 +84,10 @@ def build_and_visualize_graph(graph_data, title="Graph Visualization"):
     st.pyplot(plt)
 
 def main():
-    st.title("Integrated Three-Stage Pipeline with Flan-T5 and Graph Visualization")
+    st.title("Two-LLM Pipeline: Request & Reply with Graph Visualization")
     st.write(
-        "Enter a customer request to see it processed through three stages:\n\n"
-        "1. **Stage 1:** The request is converted into a structured graph (input graph).\n"
-        "2. **Stage 2:** The graph is used to generate a detailed response.\n"
-        "3. **Stage 3:** The response is converted into another structured graph (output graph) and visualized."
+        "Enter a customer request to see it transformed into a structured graph (input graph) "
+        "and then receive a detailed reply with its own structured graph (output graph)."
     )
     
     customer_request = st.text_area("Enter Customer Request:", height=150)
@@ -105,7 +99,7 @@ def main():
         
         # --- Stage 1: Request to Graph ---
         st.subheader("Stage 1: Customer Request → Input Graph")
-        model1_output = model1_translate_to_graph(customer_request)
+        model1_output = generate_input_graph(customer_request)
         st.code(model1_output, language="json")
         input_graph_data = try_parse_json(model1_output)
         if input_graph_data:
@@ -113,20 +107,16 @@ def main():
         else:
             st.warning("Could not parse a valid graph structure from Model 1 output.")
         
-        # --- Stage 2: Graph → Detailed Response ---
-        st.subheader("Stage 2: Input Graph → Detailed Response")
-        response_text = model2_generate_response(model1_output)
-        st.write(response_text)
-        
-        # --- Stage 3: Detailed Response → Output Graph ---
-        st.subheader("Stage 3: Detailed Response → Output Graph")
-        model3_output = model3_response_to_graph(response_text)
-        st.code(model3_output, language="json")
-        output_graph_data = try_parse_json(model3_output)
+        # --- Stage 2: Detailed Reply with Graph ---
+        st.subheader("Stage 2: Detailed Response with Structured Graph")
+        model2_output = generate_reply_and_graph(customer_request)
+        st.write("Detailed Response and Graph (raw output):")
+        st.code(model2_output, language="json")
+        output_graph_data = try_parse_json(model2_output)
         if output_graph_data:
             build_and_visualize_graph(output_graph_data, title="Output Graph (Structured Response)")
         else:
-            st.warning("Could not parse a valid graph structure from Model 3 output.")
+            st.warning("Could not parse a valid graph structure from Model 2 output.")
 
 if __name__ == "__main__":
     main()