Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import streamlit as st
+from dotenv import load_dotenv
 from langchain.schema import HumanMessage
 from langchain_groq import ChatGroq
 from langgraph.graph import StateGraph, START, END
@@ -7,18 +8,24 @@ from pydantic import BaseModel
 from langsmith import traceable
 import traceback
 
-# ✅ Load API keys
+# ✅ Load API keys securely
+load_dotenv()
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 LANGSMITH_API_KEY = os.getenv("LANGSMITH_API_KEY")
 
-# ✅ Initialize LLM (Using Groq Llama3-8B)
+# ✅ Set API keys as environment variables
+os.environ["GROQ_API_KEY"] = GROQ_API_KEY
+os.environ["LANGCHAIN_TRACING_V2"] = "true"
+os.environ["LANGCHAIN_API_KEY"] = LANGSMITH_API_KEY
+
+# ✅ Initialize Open-Source LLM (Using Groq Llama3-8B)
 llm = ChatGroq(groq_api_key=GROQ_API_KEY, model_name="llama3-8b-8192")
 
 # ✅ Define the LegalState Model
 class LegalState(BaseModel):
     original_text: str
     tone: str
-    complexity: int = 1  # Default complexity level
+    complexity: int = 1
     rewritten_text: str = None
     summary: str = None
     key_clauses: str = None
@@ -30,8 +37,7 @@ class LegalState(BaseModel):
     comparison_result: str = None
     final_report: str = None
 
-# ✅ Function to invoke LLM
-@traceable(name="Generate Response")
+# ✅ Function to invoke LLM with error handling
 def generate_response(prompt):
     try:
         response = llm.invoke([HumanMessage(content=prompt)])
@@ -39,10 +45,11 @@ def generate_response(prompt):
     except Exception as e:
         return f"❌ Error: {str(e)}"
 
-# ✅ Define Workflow Functions
+# ✅ Define Worker Functions with LangSmith Debugging
 @traceable(name="Rewrite Legal Text")
 def rewrite_text(state: LegalState):
-    prompt = f"""Rewrite this legal text in '{state.tone}' tone with complexity level {state.complexity}:\n\n{state.original_text}"""
+    prompt = f"""Rewrite this legal text in '{state.tone}' tone with complexity level {state.complexity}:
+    {state.original_text}"""
     return {"rewritten_text": generate_response(prompt)}
 
 @traceable(name="Summarize Legal Text")
@@ -55,14 +62,23 @@ def extract_clauses(state: LegalState):
 
 @traceable(name="Final Synthesis")
 def synthesizer(state: LegalState):
-    return {"final_report": f"""
-    📜 **Final Legal Report** 📜\n\n
-    **Rewritten Legal Text:** {state.rewritten_text}\n\n
-    **Summary:** {state.summary}\n\n
-    **Key Clauses:** {state.key_clauses}\n"""}
+    final_output = f"""
+    📜 **AI-Powered Legal Document Processing Report** 📜
+    
+    **🔹 Rewritten Legal Text:**  
+    **{state.rewritten_text}**
+    
+    **🔹 Summary:**  
+    {state.summary}
+    
+    **🔹 Key Clauses Identified:**  
+    {state.key_clauses}
+    """
+    return {"final_report": final_output}
 
-# ✅ Orchestrator
+# ✅ Build LangGraph Workflow
 builder = StateGraph(LegalState)
+
 builder.add_node("rewrite_text", rewrite_text)
 builder.add_node("summarize_text", summarize_text)
 builder.add_node("extract_clauses", extract_clauses)
@@ -82,22 +98,18 @@ st.title("📜 AI-Powered Legal Text Processor")
 original_text = st.text_area("Enter Legal Text")
 tone = st.radio("Select Tone", ["Formal", "Empathetic", "Neutral", "Strength-Based"])
 
-if st.button("Generate"): 
+if st.button("Rewrite Text"):
     input_data = {"original_text": original_text, "tone": tone, "complexity": 1}
     result = graph.invoke(input_data)
-    
-    st.subheader("🔹 Rewritten Text")
-    rewritten_text = st.text_area("", value=result["rewritten_text"], height=150)
-    if st.button("Regenerate"): 
-        input_data["original_text"] = rewritten_text
+    st.session_state["rewritten_text"] = result.get("rewritten_text", "")
+    st.session_state["show_regen"] = True
+
+if "show_regen" in st.session_state and st.session_state["show_regen"]:
+    if st.button("Regenerate Text"):
+        input_data = {"original_text": original_text, "tone": tone, "complexity": 1}
         result = graph.invoke(input_data)
-        st.experimental_rerun()
-    
-    st.subheader("🔹 Summary")
-    st.write(result["summary"])
-    
-    st.subheader("🔹 Key Clauses")
-    st.write(result["key_clauses"])
-    
-    st.subheader("📜 Final Report")
-    st.write(result["final_report"])
+        st.session_state["rewritten_text"] = result.get("rewritten_text", "")
+
+if "rewritten_text" in st.session_state:
+    st.subheader("Rewritten Text")
+    st.write(f"**{st.session_state['rewritten_text']}**")