Update app.py

app.py

@@ -8,21 +8,23 @@ from pydantic import BaseModel
 from langsmith import traceable
 import traceback
 
-# Load API keys from environment
+# ✅ Load API keys securely
 load_dotenv()
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 LANGSMITH_API_KEY = os.getenv("LANGSMITH_API_KEY")
+
+# ✅ 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 LLM
+# ✅ Initialize Open-Source LLM
 llm = ChatGroq(groq_api_key=GROQ_API_KEY, model_name="llama3-8b-8192")
 
-# Define State Model
+# ✅ Define the LegalState Model
 class LegalState(BaseModel):
     original_text: str
     tone: str
-    complexity: int = 1
     rewritten_text: str = None
     summary: str = None
     key_clauses: str = None
@@ -34,7 +36,7 @@ class LegalState(BaseModel):
     comparison_result: str = None
     final_report: str = None
 
-# Function to call LLM
+# ✅ Function to invoke LLM with error handling
 def generate_response(prompt):
     try:
         response = llm.invoke([HumanMessage(content=prompt)])
@@ -42,113 +44,87 @@ 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}:
+    prompt = f"""**Rewrite this legal text in '{state.tone}' tone:**
+
     {state.original_text}"""
     return {"rewritten_text": generate_response(prompt)}
 
 @traceable(name="Summarize Legal Text")
 def summarize_text(state: LegalState):
-    return {"summary": generate_response(f"Summarize this legal text:\n\n{state.rewritten_text}")}
+    return {"summary": generate_response(f"**Summary:**\n\n{state.rewritten_text}")}
 
 @traceable(name="Extract Key Clauses")
 def extract_clauses(state: LegalState):
-    return {"key_clauses": generate_response(f"Extract key legal clauses:\n\n{state.rewritten_text}")}
+    return {"key_clauses": generate_response(f"**Key Clauses:**\n\n{state.rewritten_text}")}
 
 @traceable(name="Detect Risks in Document")
 def detect_risks(state: LegalState):
-    return {"risk_analysis": generate_response(f"Analyze for risks:\n\n{state.rewritten_text}")}
+    return {"risk_analysis": generate_response(f"**Risk Analysis:**\n\n{state.rewritten_text}")}
 
 @traceable(name="Check Compliance")
 def check_compliance(state: LegalState):
-    return {"compliance_report": generate_response(f"Check legal compliance:\n\n{state.rewritten_text}")}
+    return {"compliance_report": generate_response(f"**Compliance Report:**\n\n{state.rewritten_text}")}
 
 @traceable(name="Suggest Contract Improvements")
 def suggest_improvements(state: LegalState):
-    return {"contract_suggestions": generate_response(f"Suggest improvements:\n\n{state.rewritten_text}")}
+    return {"contract_suggestions": generate_response(f"**Contract Improvement Suggestions:**\n\n{state.rewritten_text}")}
 
 @traceable(name="Generate Legal Arguments")
 def generate_arguments(state: LegalState):
-    return {"legal_arguments": generate_response(f"Generate legal arguments:\n\n{state.rewritten_text}")}
+    return {"legal_arguments": generate_response(f"**Legal Arguments:**\n\n{state.rewritten_text}")}
 
 @traceable(name="Format Legal Document")
 def format_document(state: LegalState):
-    return {"formatted_text": generate_response(f"Format this legal document:\n\n{state.rewritten_text}")}
+    return {"formatted_text": generate_response(f"**Formatted Legal Document:**\n\n{state.rewritten_text}")}
 
 @traceable(name="Compare Original vs Rewritten")
 def compare_texts(state: LegalState):
-    return {"comparison_result": generate_response(f"Compare original vs rewritten:\n\nOriginal: {state.original_text}\nRewritten: {state.rewritten_text}")}
+    prompt = f"**Comparison:**\n\nOriginal: {state.original_text}\nRewritten: {state.rewritten_text}"
+    return {"comparison_result": generate_response(prompt)}
 
 @traceable(name="Final Synthesis")
 def synthesizer(state: LegalState):
     final_output = f"""
-    📜 **Legal Document Processing Report** 📜
-
-    **Rewritten Legal Text:**  
-    {state.rewritten_text}
-
-    **Summary:**  
-    {state.summary}
-
-    **Key Clauses:**  
-    {state.key_clauses}
-
-    **Risk Analysis:**  
-    {state.risk_analysis}
-
-    **Compliance Report:**  
-    {state.compliance_report}
-
-    **Contract Suggestions:**  
-    {state.contract_suggestions}
-
-    **Legal Arguments:**  
-    {state.legal_arguments}
-
-    **Formatted Document:**  
-    {state.formatted_text}
-
-    **Comparison:**  
-    {state.comparison_result}
+    📜 **AI-Powered Legal Document Processing Report** 📜
+
+    **🔹 Rewritten Legal Text:**\n{state.rewritten_text}
+    **🔹 Summary:**\n{state.summary}
+    **🔹 Key Clauses Identified:**\n{state.key_clauses}
+    **🔹 Risk Analysis Report:**\n{state.risk_analysis}
+    **🔹 Compliance Check Results:**\n{state.compliance_report}
+    **🔹 Contract Improvement Suggestions:**\n{state.contract_suggestions}
+    **🔹 AI-Generated Legal Arguments:**\n{state.legal_arguments}
+    **🔹 Text Formatting Enhancements:**\n{state.formatted_text}
+    **🔹 Differences from Original Text:**\n{state.comparison_result}
+    **✅ AI Final Verdict:**\nThis document has been reviewed and processed with key legal enhancements.
     """
     return {"final_report": final_output}
 
-# Build LangGraph Workflow
+# ✅ 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)
-builder.add_node("detect_risks", detect_risks)
-builder.add_node("check_compliance", check_compliance)
-builder.add_node("suggest_improvements", suggest_improvements)
-builder.add_node("generate_arguments", generate_arguments)
-builder.add_node("format_document", format_document)
-builder.add_node("compare_texts", compare_texts)
-builder.add_node("synthesizer", synthesizer)
 
+# ✅ Add Nodes
+nodes = ["rewrite_text", "summarize_text", "extract_clauses", "detect_risks", "check_compliance", "suggest_improvements", "generate_arguments", "format_document", "compare_texts", "synthesizer"]
+for node in nodes:
+    builder.add_node(node, globals()[node])
+
+# ✅ Define Execution Order
 builder.add_edge(START, "rewrite_text")
-for node in ["summarize_text", "extract_clauses", "detect_risks", "check_compliance", "suggest_improvements", "generate_arguments", "format_document", "compare_texts"]:
+for node in nodes[:-1]:
     builder.add_edge("rewrite_text", node)
     builder.add_edge(node, "synthesizer")
 builder.add_edge("synthesizer", END)
 
 graph = builder.compile()
 
-# Streamlit UI
+# ✅ Streamlit UI
 st.title("📜 AI-Powered Legal Text Processor")
-text_input = st.text_area("Enter Legal Text")
-tone_selection = st.radio("Select Tone", ["Formal", "Empathetic", "Neutral", "Strength-Based"])
+original_text = st.text_area("Enter Legal Text")
+tone = st.selectbox("Select Tone", ["Formal", "Empathetic", "Neutral", "Strength-Based"])
 if st.button("Rewrite Text"):
-    output = graph.invoke({"original_text": text_input, "tone": tone_selection, "complexity": 1})
-    st.session_state["output"] = output
-    st.session_state["regenerate"] = True
-if "regenerate" in st.session_state and st.session_state["regenerate"]:
-    if st.button("Regenerate Text"):
-        output = graph.invoke({"original_text": text_input, "tone": tone_selection, "complexity": 1})
-        st.session_state["output"] = output
-if "output" in st.session_state:
-    for key, value in st.session_state["output"].items():
-        st.markdown(f"**{key.replace('_', ' ').title()}**:")
-        st.write(value)
+    result = graph.invoke({"original_text": original_text, "tone": tone})
+    for key, value in result.items():
+        st.markdown(f"**{key.replace('_', ' ').title()}:**\n{value}")