Update app.py

app.py

@@ -13,18 +13,18 @@ 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 Open-Source LLM
+# ✅ Initialize 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
     rewritten_text: str = None
     summary: str = None
     key_clauses: str = None
@@ -44,87 +44,103 @@ def generate_response(prompt):
     except Exception as e:
         return f"❌ Error: {str(e)}"
 
-# ✅ Define Worker Functions with LangSmith Debugging
+# ✅ Define Worker Functions
 @traceable(name="Rewrite Legal Text")
 def rewrite_text(state: LegalState):
-    prompt = f"""**Rewrite this legal text in '{state.tone}' tone:**
+    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")
 def summarize_text(state: LegalState):
-    return {"summary": generate_response(f"**Summary:**\n\n{state.rewritten_text}")}
+    return {"summary": generate_response(f"Summarize this legal text:\n\n{state.rewritten_text}")}
 
 @traceable(name="Extract Key Clauses")
 def extract_clauses(state: LegalState):
-    return {"key_clauses": generate_response(f"**Key Clauses:**\n\n{state.rewritten_text}")}
+    return {"key_clauses": generate_response(f"Extract key legal clauses:\n\n{state.rewritten_text}")}
 
 @traceable(name="Detect Risks in Document")
 def detect_risks(state: LegalState):
-    return {"risk_analysis": generate_response(f"**Risk Analysis:**\n\n{state.rewritten_text}")}
+    return {"risk_analysis": generate_response(f"Analyze for risks:\n\n{state.rewritten_text}")}
 
 @traceable(name="Check Compliance")
 def check_compliance(state: LegalState):
-    return {"compliance_report": generate_response(f"**Compliance Report:**\n\n{state.rewritten_text}")}
+    return {"compliance_report": generate_response(f"Check legal compliance:\n\n{state.rewritten_text}")}
 
 @traceable(name="Suggest Contract Improvements")
 def suggest_improvements(state: LegalState):
-    return {"contract_suggestions": generate_response(f"**Contract Improvement Suggestions:**\n\n{state.rewritten_text}")}
+    return {"contract_suggestions": generate_response(f"Suggest improvements:\n\n{state.rewritten_text}")}
 
 @traceable(name="Generate Legal Arguments")
 def generate_arguments(state: LegalState):
-    return {"legal_arguments": generate_response(f"**Legal Arguments:**\n\n{state.rewritten_text}")}
+    return {"legal_arguments": generate_response(f"Generate legal arguments:\n\n{state.rewritten_text}")}
 
 @traceable(name="Format Legal Document")
 def format_document(state: LegalState):
-    return {"formatted_text": generate_response(f"**Formatted Legal Document:**\n\n{state.rewritten_text}")}
+    return {"formatted_text": generate_response(f"Format this legal document:\n\n{state.rewritten_text}")}
 
 @traceable(name="Compare Original vs Rewritten")
 def compare_texts(state: LegalState):
-    prompt = f"**Comparison:**\n\nOriginal: {state.original_text}\nRewritten: {state.rewritten_text}"
+    prompt = f"Compare original vs rewritten:\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"""
-    📜 **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
 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)
 
-# ✅ 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 nodes[:-1]:
+for node in ["summarize_text", "extract_clauses", "detect_risks", "check_compliance", "suggest_improvements", "generate_arguments", "format_document", "compare_texts"]:
     builder.add_edge("rewrite_text", node)
-    builder.add_edge(node, "synthesizer")
-builder.add_edge("synthesizer", END)
 
 graph = builder.compile()
 
 # ✅ Streamlit UI
-st.title("📜 AI-Powered Legal Text Processor")
-original_text = st.text_area("Enter Legal Text")
-tone = st.selectbox("Select Tone", ["Formal", "Empathetic", "Neutral", "Strength-Based"])
-if st.button("Rewrite Text"):
-    result = graph.invoke({"original_text": original_text, "tone": tone})
-    for key, value in result.items():
-        st.markdown(f"**{key.replace('_', ' ').title()}:**\n{value}")
+def main():
+    st.title("📜 AI-Powered Legal Text Processor")
+    st.write("A smart legal document assistant powered by LLMs.")
+
+    original_text = st.text_area("Enter Legal Text:")
+    tone = st.radio("Select Tone:", ["Formal", "Empathetic", "Neutral", "Strength-Based"])
+    
+    if st.button("Process Text"):
+        input_data = {"original_text": original_text, "tone": tone, "complexity": 1}
+        result = graph.invoke(input_data)
+        
+        st.subheader("🔹 Rewritten Text")
+        st.markdown(f"**{result['rewritten_text']}**")
+
+        st.subheader("📌 Summary")
+        st.write(result['summary'])
+        
+        st.subheader("📜 Key Clauses")
+        st.write(result['key_clauses'])
+        
+        st.subheader("⚠️ Risk Analysis")
+        st.write(result['risk_analysis'])
+        
+        st.subheader("✅ Compliance Report")
+        st.write(result['compliance_report'])
+        
+        st.subheader("💡 Contract Suggestions")
+        st.write(result['contract_suggestions'])
+        
+        st.subheader("⚖️ Legal Arguments")
+        st.write(result['legal_arguments'])
+        
+        st.subheader("📝 Formatted Legal Document")
+        st.write(result['formatted_text'])
+        
+        st.subheader("🔍 Comparison Result")
+        st.write(result['comparison_result'])
+
+if __name__ == "__main__":
+    main()