add the doc_editor to chat agent

agent_api.py

@@ -36,6 +36,7 @@ from agents.pdf_analyzer import PDFAnalyzerAgent
 from agents.doc_editor import DocumentEditorAgent
 from agents.doc_assistant import DocAssistant
 from langchain_openai import ChatOpenAI
+from langchain_xai import ChatXAI
 from mistralai import Mistral
 import logging
 import traceback
@@ -92,6 +93,10 @@ class LLMConfig:
         #     api_key=os.getenv("GOOGLE_API_KEY"),
         #     thinking_level="low"
         # )
+        self.llm = ChatXAI(
+            model=os.getenv("XAI_TOOL_MODEL"),
+        )
+
         # logger.info("✅ LLMConfig initialized:")
         # logger.info(f"   - OpenAI LLM: {os.getenv('LLM_MODEL', 'gpt-5-nano-2025-08-07')}")
         # logger.info(f"   - Gemini LLM: {os.getenv('GEMINI_TOOL_MODEL', 'gemini-3-flash-preview')} (for tool calling)")

agents/chat_agent.py

@@ -123,6 +123,14 @@ class CyberLegalAgent:
                 # Inject user_id for create_draft_document tool
                 if tool_call['name'] == "create_draft_document":
                     args["user_id"] = state.get("user_id")
+                    args["instruction"] = state.get("user_query", "")
+                    # Inject conversation_history for context
+                    args["conversation_history"] = state.get("conversation_history", [])
+                    logger.info(f"💬 Injecting conversation_history for create_draft_document: {len(args['conversation_history'])} messages")
+                    # Inject doc_summaries for context
+                    args["doc_summaries"] = state.get("doc_summaries", [])
+                    logger.info(f"📄 Injecting doc_summaries for create_draft_document: {len(args['doc_summaries'])} documents")
+                    logger.info(f"📝 Using user_query as instruction for create_draft_document")
                     logger.info(f"📝 Injecting user_id for create_draft_document: {args['user_id']}")
                 
                 # Convert to implementation name

prompts/main.py

@@ -75,6 +75,12 @@ You are one of the agent in Hexiagon AI platform, a juridical platform that help
 1. **query_knowledge_graph**: Search legal documents to answer questions about law, regulations and directives.
 2. **search_web**: Search the web for current information, court decisions, or news that may not be in the knowledge graph.
 3. **retrieve_lawyer_document**: Retrieve the full content of a specific document from your document database. Use this tool when you need to inspect the detailed text, clauses, or specific provisions of a document mentioned in your documents_tree. But avoid it if the summary key details and actors are sufficient as it can introduce long document in the chat.
+4. **create_draft_document**: Create a new document draft by building it from scratch using an execution plan. This tool creates an empty document, then uses the doc_editor agent to build it step-by-step according to your plan, and saves it directly to your "My documents" section on the platform. The document is created with full access to the conversation history and relevant document summaries for context. Use this when the user asks you to create, draft, or build a new legal document (contract, agreement, letter, etc.). You will need to provide:
+   - `path_with_filename`: Full path and filename with extension (e.g., "Contracts/Contrat_bail.pdf" or "Courriers/Lettre_relance.pdf")
+   - `plan_to_build_document`: A numbered list of 1-6 steps describing how to build the document (e.g., "1. Create header with title\n2. Add section for parties\n3. Add main clauses\n4. Add signature section")
+   - `instruction`: The user's instruction for what to create (this is automatically provided from their query)
+   
+   **CRITICAL**: Before calling create_draft_document, you MUST ensure the request is sufficiently clear and detailed. If the user's request is vague or missing essential information (e.g., "create a contract" without specifying type, parties, key terms), ask clarifying questions first. Only proceed with document creation when you have enough details to create a meaningful, complete document.
 
 ### Tool-Calling Process
 You operate in an iterative loop:
@@ -89,6 +95,12 @@ You operate in an iterative loop:
 3. Analyze legal implications, potential interpretations, and relevant jurisprudence where applicable.
 4. Create a section at the end of your response called "References" that lists the source documents used to answer the user's question with full legal citations.
 5. If the user asks a specific detail about a document from your documents_tree, use the retrieve_lawyer_document tool, but avoid it if the summary key details and actors are sufficient as it can introduce long document in the chat.
+6. **CRITICAL - Before creating documents**: When the user asks you to create a document, you MUST ensure the request is sufficiently clear and detailed. If the user's request is vague or missing essential information (e.g., "create a contract" without specifying type, parties, key terms), ask clarifying questions first. Only proceed with document creation when you have enough details to create a meaningful, complete document.
+7. **After creating documents**: When you create a document using the create_draft_document tool, you MUST clearly explain to the user:
+   - What document was created (title, type)
+   - What the document contains (summary of main sections)
+   - **EXACTLY where the document was saved** in their "My documents" section (full path with filename, e.g., "The document has been saved to: Contracts/Contrat_bail.pdf in your My documents section")
+   - This information is critical so the user can find and access their document
 
 ### Tone
 - Professional and authoritative

requirements.txt

@@ -13,7 +13,7 @@ langchain-openai>=0.1.0
 langchain-community>=0.0.20
 langchain-google-genai>=1.0.0
 mistralai>=1.0.0
-
+langchain-xai==1.2.2
 # FastAPI and server dependencies
 fastapi>=0.104.0
 uvicorn[standard]>=0.24.0

tests/test_create_draft_document.py

@@ -20,9 +20,9 @@ async def test_tool_facade():
     
     # This should return immediately (facade just returns)
     result = await create_draft_document.ainvoke({
-        "title": "Test Document",
-        "content": "<h1>Test</h1><p>This is a test document</p>",
-        "path": "Test/"
+        "path_with_filename": "Test/Test_Document.pdf",
+        "plan_to_build_document": "1. Create title\n2. Add content",
+        "instruction": "Create a test document"
     })
     
     print(f"✅ Facade test passed: {result}")
@@ -36,38 +36,46 @@ async def test_tool_real():
     # Mock user_id
     test_user_id = os.getenv("TEST_USER_ID", "test-user-uuid-123")
     
-    # Test with various paths
+    # Test with various paths and plans
     test_cases = [
         {
-            "title": "Contract de bail",
-            "content": "<h1>Contrat de bail</h1><p>Ce contrat est conclu entre...</p>",
-            "path": "Contracts/",
-            "expected_path": "./Contracts/Contract de bail.pdf"
+            "path_with_filename": "Contracts/Contract_de_bail.pdf",
+            "plan_to_build_document": """1. Create header with title "Contrat de bail"
+2. Add section for parties (bailleur and preneur)
+3. Add main clauses for the lease
+4. Add signature section""",
+            "instruction": "Create a commercial lease contract",
+            "expected_path": "Contracts/Contract_de_bail.pdf"
         },
         {
-            "title": "Mon document",
-            "content": "<h1>Document</h1><p>Contenu...</p>",
-            "path": "",
-            "expected_path": "./Mon document.pdf"
+            "path_with_filename": "Mon_document.pdf",
+            "plan_to_build_document": """1. Add document title
+2. Add introductory paragraph
+3. Add main content sections""",
+            "instruction": "Create a simple document",
+            "expected_path": "Mon_document.pdf"
         },
         {
-            "title": "Note juridique",
-            "content": "<h1>Note</h1><p>Contenu juridique...</p>",
-            "path": "Drafts/Legal/",
-            "expected_path": "./Drafts/Legal/Note juridique.pdf"
+            "path_with_filename": "Drafts/Legal/Note_juridique.pdf",
+            "plan_to_build_document": """1. Create title "Note juridique"
+2. Add legal context
+3. Add analysis section
+4. Add conclusions""",
+            "instruction": "Create a legal note",
+            "expected_path": "Drafts/Legal/Note_juridique.pdf"
         }
     ]
     
     for i, test_case in enumerate(test_cases, 1):
-        print(f"\n  Test case {i}: {test_case['title']}")
-        print(f"  Path: {test_case['path']}")
+        print(f"\n  Test case {i}: {test_case['path_with_filename']}")
+        print(f"  Instruction: {test_case['instruction']}")
         print(f"  Expected: {test_case['expected_path']}")
         
         result = await _create_draft_document.ainvoke({
             "user_id": test_user_id,
-            "title": test_case["title"],
-            "content": test_case["content"],
-            "path": test_case["path"]
+            "path_with_filename": test_case["path_with_filename"],
+            "plan_to_build_document": test_case["plan_to_build_document"],
+            "instruction": test_case["instruction"]
         })
         
         print(f"  Result: {result}")
@@ -82,39 +90,41 @@ async def test_tool_real():
     return True
 
 
-async def test_path_normalization():
-    """Test path normalization logic"""
-    print("\n🧪 Testing path normalization...")
+async def test_path_parsing():
+    """Test path parsing logic"""
+    print("\n🧪 Testing path parsing...")
     
     test_cases = [
-        ("Contracts/", "./Contracts/"),
-        ("Contracts", "./Contracts/"),
-        ("", "./"),
-        ("./Contracts/", "./Contracts/"),
-        ("Drafts/Legal/", "./Drafts/Legal/"),
-        ("./Drafts/Legal", "./Drafts/Legal/")
+        ("Contracts/Test.pdf", "Contracts", "Test.pdf", "Test"),
+        ("Test.pdf", "", "Test.pdf", "Test"),
+        ("Drafts/Legal/Note.pdf", "Drafts/Legal", "Note.pdf", "Note"),
+        ("./Contracts/Test.pdf", "./Contracts", "Test.pdf", "Test"),
     ]
     
-    for input_path, expected_suffix in test_cases:
-        # Simulate the path normalization logic
-        path = input_path
-        if path:
-            if path.startswith('./'):
-                path = path[2:]
-            if not path.endswith('/'):
-                path += '/'
+    for path_with_filename, expected_path, expected_filename, expected_title in test_cases:
+        # Simulate the path parsing logic
+        parts = path_with_filename.split('/')
+        if len(parts) > 1:
+            path = '/'.join(parts[:-1])
+            filename = parts[-1]
         else:
             path = ''
+            filename = parts[0]
         
-        full_path = f"./{path}Test.pdf"
+        # Remove .pdf extension
+        title = filename.replace('.pdf', '')
         
-        print(f"  Input: '{input_path}' → Output: '{full_path}'")
-        if full_path.startswith(expected_suffix):
+        print(f"  Input: '{path_with_filename}'")
+        print(f"    → Path: '{path}'")
+        print(f"    → Filename: '{filename}'")
+        print(f"    → Title: '{title}'")
+        
+        if path == expected_path and filename == expected_filename and title == expected_title:
             print(f"  ✅ Correct")
         else:
-            print(f"  ⚠️  Expected to start with: {expected_suffix}")
+            print(f"  ⚠️  Mismatch - Expected: path='{expected_path}', filename='{expected_filename}', title='{expected_title}'")
     
-    print("\n✅ Path normalization test completed")
+    print("\n✅ Path parsing test completed")
     return True
 
 
@@ -131,8 +141,8 @@ async def main():
         # Test real implementation
         await test_tool_real()
         
-        # Test path normalization
-        await test_path_normalization()
+        # Test path parsing
+        await test_path_parsing()
         
         print("\n" + "=" * 80)
         print("✅ All tests completed successfully!")

utils/tools.py

@@ -363,59 +363,192 @@ async def _retrieve_lawyer_document(
 
 @tool
 async def create_draft_document(
-    title: str,
-    content: str,
-    path: str
+    path_with_filename: str,
+    plan_to_build_document: str,
+    instruction: str
 ) -> str:
     """
-    Create a new document draft and save it to "My documents" via Supabase.
+    Create a new document draft by building it with doc_editor and save to "My documents" via Supabase.
     
-    This tool saves an HTML document as a PDF draft in the user's document storage.
-    The document will be stored in the specified folder path within "My Documents".
+    This tool creates an empty document, then uses doc_editor to build it according to a plan,
+    and finally saves the complete document to the user's document storage.
     
-    Use this tool when the user wants to:
-    - Create a new document draft
-    - Save a generated document
-    - Store a document in their document library
+    Use this tool when you want to:
+    - Create a new document from scratch
+    - Build a document using an execution plan
+    - Leverage the doc_editor's iterative editing capabilities
     
     Args:
-        title: Document title (e.g., "Contract de bail", "Note juridique")
-        content: Document content in HTML format (e.g., "<h1>Title</h1><p>Content...</p>")
-        path: Folder path where to save the document
-               - Empty string "" → root folder of My Documents
-               - "Contracts/" → ./Contracts/title.pdf
-               - "Drafts/Legal/" → ./Drafts/Legal/title.pdf
+        path_with_filename: Full path and filename with extension (e.g., "Contracts/Contrat_bail.pdf")
+        plan_to_build_document: A numbered list of steps to build the document (1-6 steps recommended)
+        instruction: The user's instruction for what to create (e.g., "Create a commercial lease contract")
     
     Returns:
-        Confirmation message with document path and success status
+        Confirmation message with document path and summary of what was created
     
     Example:
         create_draft_document(
-            title="Contract de bail",
-            content="<h1>Contrat de bail</h1><p>Ce contrat est conclu entre...</p>",
-            path="Contracts/"
+            path_with_filename="Contracts/Contrat_bail.pdf",
+            plan_to_build_document="1. Create section for parties\\n2. Add main clauses\\n3. Add signature section",
+            instruction="Create a commercial lease contract"
         )
-        → Saves as "./Contracts/Contract de bail.pdf" and returns confirmation
+        → Creates document via doc_editor and saves as "./Contracts/Contrat_bail.pdf"
     """
     return
 
 @tool
 async def _create_draft_document(
+    user_id: str,
+    path_with_filename: str,
+    plan_to_build_document: str,
+    instruction: str,
+    doc_summaries: list = [],
+    conversation_history: list = []
+) -> str:
+    """
+    Create a document by building it with doc_editor and save to Supabase.
+    
+    Args:
+        user_id: User UUID (injected by the agent)
+        path_with_filename: Full path and filename (e.g., "Contracts/Contrat_bail.pdf")
+        plan_to_build_document: Execution plan for building the document
+        instruction: User's instruction for what to create
+        doc_summaries: Optional list of document summaries for context (injected from state)
+        conversation_history: Optional conversation history for context (injected from state)
+    
+    Returns:
+        Success/failure message with document path and summary
+    """
+    try:
+        # Parse path and filename
+        parts = path_with_filename.split('/')
+        if len(parts) > 1:
+            path = '/'.join(parts[:-1])
+            filename = parts[-1]
+        else:
+            path = ''
+            filename = parts[0]
+        
+        # Remove .pdf extension if present
+        title = filename.replace('.pdf', '')
+        
+        # 1. Create empty HTML document
+        empty_html = """<!DOCTYPE html>
+<html>
+<head>
+    <meta charset="UTF-8">
+    <title>Document</title>
+</head>
+<body>
+</body>
+</html>"""
+        
+        # 2. Save empty document and get document_id
+        logger.info("=" * 80)
+        logger.info("📄 STEP 1: Creating empty document on Supabase")
+        logger.info("=" * 80)
+        logger.info(f"   User ID: {user_id}")
+        logger.info(f"   Title: {title}")
+        logger.info(f"   Path: {path if path else 'root'}")
+        logger.info(f"   Empty HTML size: {len(empty_html)} bytes")
+        
+        msg, document_id = await _create_draft_document_impl(
+            user_id=user_id,
+            title=title,
+            content=empty_html,
+            path=path
+        )
+        
+        if not document_id:
+            logger.error(f"❌ Failed to create empty document: {msg}")
+            return f"❌ Failed to create document: {msg}"
+        
+        logger.info(f"✅ Empty document created successfully!")
+        logger.info(f"   Document ID: {document_id}")
+        logger.info(f"   Full path: {msg}")
+        logger.info("")
+        
+        # 3. Launch doc_editor to build the document with real-time updates
+        logger.info("=" * 80)
+        logger.info("🔧 STEP 2: Launching doc_editor to build document")
+        logger.info("=" * 80)
+        logger.info(f"   Document title: {title}")
+        logger.info(f"   Document ID: {document_id}")
+        logger.info(f"   User ID: {user_id}")
+        logger.info(f"   Instruction: {instruction}")
+        logger.info(f"   Plan length: {len(plan_to_build_document)} characters")
+        logger.info(f"   Max iterations: 10")
+        logger.info(f"   Real-time updates: ENABLED (will push to Supabase after each edit)")
+        logger.info("")
+        
+        editor_result = await doc_editor_agent.edit_document(
+            doc_text=empty_html,
+            user_instruction=instruction,
+            plan=plan_to_build_document,
+            doc_summaries=doc_summaries,
+            conversation_history=conversation_history,
+            max_iterations=10,
+            document_id=document_id,  # Pass document_id for real-time updates
+            user_id=user_id            # Pass user_id for authentication
+        )
+        
+        # 4. Return result (no need to save again - doc_editor handled real-time updates)
+        if editor_result.get('success'):
+            summary = editor_result.get('final_summary', 'No summary available')
+            iteration_count = editor_result.get('iteration_count', 0)
+            final_doc_size = len(editor_result.get('doc_text', ''))
+            
+            logger.info("=" * 80)
+            logger.info("✅ STEP 3: Document creation completed successfully!")
+            logger.info("=" * 80)
+            logger.info(f"   Document path: {path_with_filename}")
+            logger.info(f"   Document ID: {document_id}")
+            logger.info(f"   Iterations completed: {iteration_count}")
+            logger.info(f"   Final document size: {final_doc_size} bytes")
+            logger.info(f"   Summary: {summary}")
+            logger.info("")
+            logger.info("📋 Document created via real-time updates:")
+            logger.info("   - Empty document created on Supabase")
+            logger.info("   - doc_editor built document iteratively")
+            logger.info("   - Each edit pushed to Supabase in real-time")
+            logger.info("   - No final save needed (already up-to-date)")
+            logger.info("=" * 80)
+            
+            return f"✅ Document created: {path_with_filename}\n\n{summary}"
+        else:
+            logger.error("=" * 80)
+            logger.error("❌ STEP 3: Document creation failed!")
+            logger.error("=" * 80)
+            logger.error(f"   Document path: {path_with_filename}")
+            logger.error(f"   Document ID: {document_id}")
+            logger.error(f"   Error message: {editor_result.get('message', 'Unknown error')}")
+            logger.error(f"   Success flag: {editor_result.get('success')}")
+            logger.error("=" * 80)
+            
+            return f"❌ Document creation failed: {editor_result.get('message', 'Unknown error')}"
+                
+    except Exception as e:
+        logger.error(f"Error in _create_draft_document: {str(e)}")
+        return f"❌ Error creating document: {str(e)}"
+
+
+async def _create_draft_document_impl(
     user_id: str,
     title: str,
     content: str,
     path: str
-) -> str:
+) -> tuple[str, str]:
     """
-    Real implementation of create_draft_document - calls Supabase endpoint.
+    Helper function to save a document to Supabase (implementation).
+    
     Args:
-        user_id: User UUID (injected by the agent)
+        user_id: User UUID
         title: Document title
         content: Document HTML content
         path: Folder path
     
     Returns:
-        Success/failure message with document path
+        tuple: (message, document_id) - Success/failure message and document_id (or None on error)
     """
     try:
         import httpx
@@ -436,6 +569,7 @@ async def _create_draft_document(
             
         full_path = f"./{path}{title}.pdf"
         
+        # URL already has /functions/v1/ (correct)
         endpoint_url = f"{base_url}/functions/v1/create-document-from-html"
         
         request_body = {
@@ -455,25 +589,27 @@ async def _create_draft_document(
             )
             
             if response.status_code == 200:
-                return f"✅ Document successfully saved to: {full_path}"
+                data = response.json()
+                document_id = data.get('document_id')
+                return f"✅ Document successfully saved to: {full_path}", document_id
             elif response.status_code == 400:
                 error_data = response.json() if response.headers.get('content-type', '').startswith('application/json') else {}
-                return f"❌ Bad request: {error_data.get('error', 'Unknown error')}"
+                return f"❌ Bad request: {error_data.get('error', 'Unknown error')}", None
             elif response.status_code == 401:
-                return "❌ Authentication failed: Invalid API key"
+                return "❌ Authentication failed: Invalid API key", None
             elif response.status_code == 403:
-                return "❌ Access denied: You do not have permission to save documents"
+                return "❌ Access denied: You do not have permission to save documents", None
             elif response.status_code == 500:
-                return "❌ Server error: Failed to save document"
+                return "❌ Server error: Failed to save document", None
             else:
-                return f"❌ Error: HTTP {response.status_code} - {response.text}"
+                return f"❌ Error: HTTP {response.status_code} - {response.text}", None
                 
     except httpx.TimeoutError:
-        return "❌ Error: Timeout while saving document"
+        return "❌ Error: Timeout while saving document", None
     except httpx.RequestError as e:
-        return f"❌ Error: Failed to connect to document server: {str(e)}"
+        return f"❌ Error: Failed to connect to document server: {str(e)}", None
     except Exception as e:
-        return f"❌ Error saving document: {str(e)}"
+        return f"❌ Error saving document: {str(e)}", None
 
 
 # ============ DOC ASSISTANT TOOLS ============
@@ -941,8 +1077,8 @@ async def _attempt_completion(message: str) -> Dict[str, Any]:
 
 
 # Export tool sets for different user types
-tools_for_client_facade = [query_knowledge_graph, find_lawyers, message_lawyer, search_web, create_draft_document]
-tools_for_client = [_query_knowledge_graph, _find_lawyers, _message_lawyer, search_web, _create_draft_document]
+tools_for_client_facade = [query_knowledge_graph, find_lawyers, message_lawyer, search_web]
+tools_for_client = [_query_knowledge_graph, _find_lawyers, _message_lawyer, search_web]
 tools_for_lawyer_facade = [query_knowledge_graph, search_web, retrieve_lawyer_document, create_draft_document]
 tools_for_lawyer = [_query_knowledge_graph, search_web, _retrieve_lawyer_document, _create_draft_document]
 
@@ -953,4 +1089,4 @@ tools_for_doc_assistant = [_query_knowledge_graph, _retrieve_lawyer_document, _e
 tools_for_doc_editor_facade = [replace_html, add_html, delete_html, view_current_document, attempt_completion]
 tools_for_doc_editor = [_replace_html, _add_html, _delete_html, _view_current_document, _attempt_completion]
 
-tools = tools_for_client
+tools = tools_for_client