Spaces:

Cyberlgl
/

CyberLegalAIendpoint

Running

File size: 7,366 Bytes

#!/usr/bin/env python3
"""
Test script for the Document Editor Agent
Demonstrates Cline-like document editing with HTML
"""

import asyncio
import os
import sys
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Add parent directory to path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from langchain_openai import ChatOpenAI
from agents.doc_editor import DocumentEditorAgent


# Example HTML document - Service Contract
SAMPLE_CONTRACT = """<!DOCTYPE html>
<html>
<head>
    <title>Contrat de prestation de services</title>
</head>
<body>
    <h1>CONTRAT DE PRESTATION DE SERVICES</h1>
    
    <h2>Article 1 - Objet</h2>
    <p>Le présent contrat a pour objet de définir les conditions dans lesquelles <strong>la Société X</strong> (ci-après « le Prestataire ») s'engage à fournir des services à <strong>la Société Y</strong> (ci-après « le Client »).</p>
    
    <h2>Article 2 - Durée</h2>
    <p>Le contrat prend effet le <em>1er janvier 2026</em> pour une durée de <u>12 mois</u>.</p>
    
    <ul>
        <li>Renouvellement tacite</li>
        <li>Préavis de 3 mois</li>
    </ul>
    
    <h2>Article 3 - Tarification</h2>
    <p>Les services sont facturés selon les conditions suivantes :</p>
    <table>
        <thead>
            <tr>
                <th>Service</th>
                <th>Tarif</th>
            </tr>
        </thead>
        <tbody>
            <tr>
                <td>Consultation</td>
                <td>200€/heure</td>
            </tr>
            <tr>
                <td>Audit</td>
                <td>500€/jour</td>
            </tr>
        </tbody>
    </table>
</body>
</html>"""


async def test_document_editor():
    """Test document editor agent with sample contract."""
    
    # Initialize LLM (requires OPENAI_API_KEY environment variable)
    api_key = os.getenv("OPENAI_API_KEY") or os.getenv("CEREBRAS_API_KEY")
    if not api_key:
        print("❌ Error: OPENAI_API_KEY or CEREBRAS_API_KEY environment variable not set")
        return
    
    llm = ChatOpenAI(
        model=os.getenv("LLM_MODEL", "gpt-4o-mini"),
        temperature=0,
        api_key=api_key,
        base_url=os.getenv("OPENAI_API_BASE", None)
    )
    
    # Initialize document editor agent
    doc_editor = DocumentEditorAgent(llm=llm, llm_tool_calling=llm)
    
    print("=" * 80)
    print("📄 ORIGINAL DOCUMENT (HTML)")
    print("=" * 80)
    print(SAMPLE_CONTRACT)
    print("\n")
    
    # Test 1: Simple text replacement
    print("=" * 80)
    print("🔧 TEST 1: Change contract duration from 12 to 24 months")
    print("=" * 80)
    
    result1 = await doc_editor.edit_document(
        doc_text=SAMPLE_CONTRACT,
        user_instruction="Change the contract duration from 12 months to 24 months in Article 2",
        doc_summaries=["Service contract between Company X and Company Y"],
        max_iterations=5
    )
    
    print(f"✅ Success: {result1['success']}")
    print(f"📝 Message: {result1['message']}")
    print(f"🔄 Iterations: {result1['iteration_count']}")
    
    if result1['success']:
        print("\n📄 MODIFIED DOCUMENT:")
        print(result1['doc_text'])
        
        # Update for next test
        current_doc = result1['doc_text']
    else:
        current_doc = SAMPLE_CONTRACT
    
    print("\n" + "=" * 80)
    print("🔧 TEST 2: Add Article 4 about confidentiality")
    print("=" * 80)
    
    result2 = await doc_editor.edit_document(
        doc_text=current_doc,
        user_instruction="Add a new Article 4 after Article 3 that covers confidentiality with a heading 'Article 4 - Confidentialité' and a paragraph explaining that all information exchanged remains confidential",
        doc_summaries=["Service contract between Company X and Company Y"],
        max_iterations=8
    )
    
    print(f"✅ Success: {result2['success']}")
    print(f"📝 Message: {result2['message']}")
    print(f"🔄 Iterations: {result2['iteration_count']}")
    
    if result2['success']:
        print("\n📄 MODIFIED DOCUMENT:")
        print(result2['doc_text'])
    
    print("\n" + "=" * 80)
    print("✅ TESTS COMPLETED")
    print("=" * 80)


async def test_tools_directly():
    """Test document editor tools directly without the agent."""
    from utils.editor_tools import replace_html, add_html, delete_html, view_current_document, attempt_completion
    
    print("=" * 80)
    print("🔧 TESTING TOOLS DIRECTLY")
    print("=" * 80)
    
    current_doc = SAMPLE_CONTRACT
    
    # Test 1: replace_html
    print("\n1️⃣ Testing 'replace_html' tool...")
    result = await replace_html.ainvoke({
        "doc_text": current_doc,
        "search": "<u>12 mois</u>",
        "replace": "<u>24 mois</u>",
        "expected_matches": 1
    })
    
    if result['ok']:
        print(f"✅ Replace successful! Found {result['matches']} matches")
        current_doc = result['doc_text']
    else:
        print(f"❌ Replace failed: {result['error']}")
    
    # Test 2: add_html
    print("\n2️⃣ Testing 'add_html' tool...")
    result = await add_html.ainvoke({
        "doc_text": current_doc,
        "anchor_search": "<h2>Article 1 - Objet</h2>",
        "insert": "<p><strong>Article ajouté : </strong>Ceci est un nouveau paragraphe inséré avant Article 1.</p>",
        "position": "before",
        "expected_matches": 1
    })
    
    if result['ok']:
        print(f"✅ Add successful! Inserted before anchor")
        current_doc = result['doc_text']
    else:
        print(f"❌ Add failed: {result['error']}")
    
    # Test 3: delete_html
    print("\n3️⃣ Testing 'delete_html' tool...")
    result = await delete_html.ainvoke({
        "doc_text": current_doc,
        "search": "<li>Renouvellement tacite</li>",
        "expected_matches": 1
    })
    
    if result['ok']:
        print(f"✅ Delete successful! Removed {result['matches']} element")
        current_doc = result['doc_text']
    else:
        print(f"❌ Delete failed: {result['error']}")
    
    print("\n4️⃣ Testing 'view_current_document' tool...")
    # Note: The tool now has no parameters - document is injected by the workflow
    result = await view_current_document.ainvoke()
    if result['ok']:
        print(f"✅ view_current_document works!")
    else:
        print(f"❌ view_current_document failed")
    
    print("\n5️⃣ Final document after tool tests:")
    print(current_doc)
    
    # Test 6: attempt_completion
    print("\n6️⃣ Testing 'attempt_completion' tool...")
    result = await attempt_completion.ainvoke({
        "message": "Successfully modified the HTML document"
    })
    print(f"✅ Completion: {result}")


if __name__ == "__main__":
    print("🚀 Document Editor Agent Test Suite\n")
    
    # Test tools directly first (simpler, no LLM needed)
    print("\n" + "=" * 80)
    print("PART 1: Testing Tools Directly")
    print("=" * 80)
    asyncio.run(test_tools_directly())
    
    # Test full agent workflow (requires LLM)
    print("\n\n" + "=" * 80)
    print("PART 2: Testing Full Agent Workflow (requires OPENAI_API_KEY or CEREBRAS_API_KEY)")
    print("=" * 80)
    print("Note: This part requires API key environment variable")
    print("=" * 80 + "\n")
    
    asyncio.run(test_document_editor())