Spaces:
Running
Running
| #!/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()) |