Update app.py

app.py

@@ -1,35 +1,155 @@
-# Update the `assistant` function to extract only the rewritten text without extra explanation
-# Inside your `assistant` function, update the part that handles rewriting and regenerating:
-
-# Example for the rewriting case
-if "Rewrite this text in" in latest_msg:
-    tone = latest_msg.split(" in ")[1].split(" tone:")[0]
-    text = latest_msg.split("tone: ")[1]
-    prompt = f"Rewrite the following legal text related to special education in California in a {tone} tone, ensuring it remains compliant with state and federal regulations. Only return the rewritten text without explanation or bullet points: {text}"
-    response = llm.invoke(prompt)
-    rewritten_text = response.content
-    return {"messages": [AIMessage(content=rewritten_text)]}
-
-# Similarly update the regenerate block
-elif latest_msg == "regenerate":
-    original_text = get_original_text(messages)
-    if not original_text:
-        return {"messages": [AIMessage(content="Error: Original text not found.")]}
-    _, current_tone = get_current_text_and_tone(messages)
-    if not current_tone:
-        return {"messages": [AIMessage(content="Error: Current tone not found.")]}
-    prompt = f"Rewrite the following legal text related to special education in California in a {current_tone} tone, ensuring it remains compliant with state and federal regulations. Only return the rewritten text without explanation or bullet points: {original_text}"
-    response = llm.invoke(prompt)
-    new_text = response.content
-    return {"messages": [AIMessage(content=new_text)]}
-
-# And the feedback block
-elif latest_msg.startswith("feedback:"):
-    feedback = latest_msg.split("feedback: ")[1].strip()
-    current_text, current_tone = get_current_text_and_tone(messages)
-    if not current_text or not current_tone:
-        return {"messages": [AIMessage(content="Error: Current text or tone not found.")]}
-    prompt = f"Refine the following legal text related to special education in California based on this feedback: '{feedback}'. Maintain the {current_tone} tone and ensure compliance with state and federal regulations. Only return the rewritten text without explanation or bullet points. Text: {current_text}"
-    response = llm.invoke(prompt)
-    refined_text = response.content
-    return {"messages": [AIMessage(content=refined_text)]}
+import os
+import uuid
+import gradio as gr
+from langgraph.graph import MessagesState, StateGraph, START
+from langgraph.checkpoint.memory import MemorySaver
+from langchain_groq import ChatGroq
+from langchain_core.messages import HumanMessage, AIMessage
+
+# Get API key from environment variables
+GROQ_API_KEY = os.environ['GROQ_API_KEY']
+
+# Initialize the Groq LLM
+llm = ChatGroq(groq_api_key=GROQ_API_KEY, model_name="llama3-8b-8192")
+
+# Helper functions
+def get_original_text(messages):
+    for msg in messages:
+        if isinstance(msg, HumanMessage) and "Rewrite this text in" in msg.content:
+            return msg.content.split("tone: ")[1]
+    return ""
+
+def get_current_text_and_tone(messages):
+    for msg in reversed(messages):
+        if isinstance(msg, AIMessage):
+            return msg.content, "Unknown"  # Simplified since we now strip metadata
+    return None, None
+
+def extract_display_text(message):
+    return message.strip()
+
+# Assistant node with clean output
+def assistant(state: MessagesState):
+    messages = state["messages"]
+    latest_msg = messages[-1].content
+
+    if "Rewrite this text in" in latest_msg:
+        tone = latest_msg.split(" in ")[1].split(" tone:")[0]
+        text = latest_msg.split("tone: ")[1]
+        prompt = f"Rewrite the following legal text related to special education in California in a {tone} tone, ensuring it remains compliant with state and federal regulations. Only return the rewritten text, without any explanation or prefix.\n\nText: {text}"
+        response = llm.invoke(prompt)
+        rewritten_text = response.content.strip()
+        return {"messages": [AIMessage(content=rewritten_text)]}
+
+    elif latest_msg == "regenerate":
+        original_text = get_original_text(messages)
+        if not original_text:
+            return {"messages": [AIMessage(content="Error: Original text not found.")]}
+        _, current_tone = get_current_text_and_tone(messages)
+        if not current_tone:
+            return {"messages": [AIMessage(content="Error: Current tone not found.")]}
+        prompt = f"Rewrite the following legal text related to special education in California in a {current_tone} tone. Only return the rewritten text, without any explanation or prefix.\n\nText: {original_text}"
+        response = llm.invoke(prompt)
+        new_text = response.content.strip()
+        return {"messages": [AIMessage(content=new_text)]}
+
+    elif latest_msg.startswith("feedback:"):
+        feedback = latest_msg.split("feedback: ")[1].strip()
+        current_text, current_tone = get_current_text_and_tone(messages)
+        if not current_text or not current_tone:
+            return {"messages": [AIMessage(content="Error: Current text or tone not found.")]}
+        prompt = f"Refine the following legal text related to special education in California based on this feedback: '{feedback}'. Maintain the {current_tone} tone and ensure compliance with regulations. Only return the refined text, without any explanation or prefix.\n\nText: {current_text}"
+        response = llm.invoke(prompt)
+        refined_text = response.content.strip()
+        return {"messages": [AIMessage(content=refined_text)]}
+
+    elif latest_msg.lower() == "approve":
+        current_text, _ = get_current_text_and_tone(messages)
+        if not current_text:
+            return {"messages": [AIMessage(content="Error: No text to approve.")]}
+        return {"messages": [AIMessage(content=f"Text approved: {current_text}")]}
+    else:
+        return {"messages": [AIMessage(content="Invalid command.")]}
+
+# Human feedback node
+def human_feedback(state: MessagesState):
+    pass
+
+# Build LangGraph
+builder = StateGraph(MessagesState)
+builder.add_node("assistant", assistant)
+builder.add_node("human_feedback", human_feedback)
+builder.add_edge(START, "human_feedback")
+builder.add_edge("human_feedback", "assistant")
+builder.add_edge("assistant", "human_feedback")
+memory = MemorySaver()
+graph = builder.compile(interrupt_before=["human_feedback"], checkpointer=memory)
+
+# Gradio logic
+def initial_submit_fn(legal_text, initial_tone, state):
+    if state["configurable"]["thread_id"] == "1":
+        state["configurable"]["thread_id"] = str(uuid.uuid4())
+    if not legal_text.strip():
+        return "Please enter legal text.", state
+    command = f"Rewrite this text in {initial_tone} tone: {legal_text}"
+    graph.update_state(state, {"messages": [HumanMessage(content=command)]}, as_node="human_feedback")
+    for event in graph.stream(None, state, stream_mode="values"):
+        if "messages" in event:
+            rewritten_text = event["messages"][-1].content
+            display_text = extract_display_text(rewritten_text)
+    return display_text, state
+
+def regenerate_fn(state):
+    graph.update_state(state, {"messages": [HumanMessage(content="regenerate")]}, as_node="human_feedback")
+    for event in graph.stream(None, state, stream_mode="values"):
+        if "messages" in event:
+            rewritten_text = event["messages"][-1].content
+            display_text = extract_display_text(rewritten_text)
+    return display_text, state
+
+def submit_feedback_fn(feedback_text, state):
+    if not feedback_text.strip():
+        return "Please enter feedback.", state, ""
+    command = f"feedback: {feedback_text}"
+    graph.update_state(state, {"messages": [HumanMessage(content=command)]}, as_node="human_feedback")
+    for event in graph.stream(None, state, stream_mode="values"):
+        if "messages" in event:
+            rewritten_text = event["messages"][-1].content
+            display_text = extract_display_text(rewritten_text)
+    return display_text, state, ""
+
+def approve_fn(state):
+    graph.update_state(state, {"messages": [HumanMessage(content="approve")]}, as_node="human_feedback")
+    for event in graph.stream(None, state, stream_mode="values"):
+        if "messages" in event:
+            confirmation = event["messages"][-1].content
+    return confirmation, state
+
+# Gradio UI
+with gr.Blocks(title="Legal Text Rewriter for Special Education in California") as demo:
+    state = gr.State(value={"configurable": {"thread_id": "1"}})
+
+    with gr.Column():
+        legal_text = gr.Textbox(label="Legal Text", placeholder="Enter your legal text here")
+        initial_tone = gr.Dropdown(
+            label="Initial Tone",
+            choices=["Formal", "Empathetic", "Neutral", "Strength-Based"],
+            value="Formal"
+        )
+        submit_btn = gr.Button("Submit")
+
+        rewritten_text_display = gr.Textbox(label="Rewritten Text", interactive=False)
+        gr.Markdown("")
+
+        regenerate_btn = gr.Button("Regenerate")
+        feedback_textbox = gr.Textbox(label="Feedback", placeholder="Enter feedback to refine the text")
+        submit_feedback_btn = gr.Button("Submit Feedback")
+        approve_btn = gr.Button("Approve")
+
+    # Button bindings
+    submit_btn.click(initial_submit_fn, [legal_text, initial_tone, state], [rewritten_text_display, state])
+    regenerate_btn.click(regenerate_fn, [state], [rewritten_text_display, state])
+    submit_feedback_btn.click(submit_feedback_fn, [feedback_textbox, state], [rewritten_text_display, state, feedback_textbox])
+    approve_btn.click(approve_fn, [state], [rewritten_text_display, state])
+
+demo.launch()