| """Test the classification agent with sample contract clauses.""" |
|
|
| import json |
| import os |
| from agents.classification_agent import classify_clause, classification_node |
|
|
| OUTPUT_DIR = os.path.join(os.path.dirname(__file__), "..", "outputs") |
|
|
| SAMPLE_CLAUSES = [ |
| { |
| "id": 1, |
| "text": "This Agreement shall be governed by and construed in accordance with the laws of the State of Delaware.", |
| "section": "General", |
| }, |
| { |
| "id": 2, |
| "text": "During the term of this Agreement and for a period of two years thereafter, the Employee shall not directly or indirectly engage in any business that competes with the Company.", |
| "section": "Restrictive Covenants", |
| }, |
| { |
| "id": 3, |
| "text": "Either party may terminate this Agreement at any time upon thirty (30) days prior written notice to the other party.", |
| "section": "Termination", |
| }, |
| { |
| "id": 4, |
| "text": "The Company shall indemnify and hold harmless the Contractor from any claims, damages, or expenses arising from the Company's breach of this Agreement.", |
| "section": "Liability", |
| }, |
| { |
| "id": 5, |
| "text": "The Licensor grants to the Licensee a non-exclusive, non-transferable license to use the Software solely for internal business purposes.", |
| "section": "License", |
| }, |
| ] |
|
|
|
|
| def test_single_clause(): |
| """Test classifying a single clause.""" |
| result = classify_clause(SAMPLE_CLAUSES[0]["text"]) |
| print(f"Clause: {SAMPLE_CLAUSES[0]['text'][:80]}...") |
| print(f"Type: {result['clause_type']}") |
| print(f"Conf: {result['confidence']}") |
| print(f"Why: {result['reasoning']}") |
| print() |
|
|
|
|
| def test_classification_node(): |
| """Test the full LangGraph node with multiple clauses.""" |
| state = {"clauses": SAMPLE_CLAUSES} |
| result = classification_node(state) |
|
|
| print(f"{'ID':<4} {'Classified As':<35} {'Confidence':<12} {'Section'}") |
| print("-" * 80) |
| for clause in result["classified_clauses"]: |
| print(f"{clause['id']:<4} {clause['clause_type']:<35} {clause['confidence']:<12.2f} {clause['section']}") |
|
|
| |
| output_path = os.path.join(OUTPUT_DIR, "classification_results.json") |
| with open(output_path, "w") as f: |
| json.dump(result["classified_clauses"], f, indent=2) |
| print(f"\nResults saved to {output_path}") |
|
|
|
|
| if __name__ == "__main__": |
| print("Single Clause Test \n") |
| test_single_clause() |
|
|
| print("Full Node Test \n") |
| test_classification_node() |
|
|
