demo_with_doc.py

#!/usr/bin/env python3
"""
Complete Demo: Create KB -> Upload Doc -> Process -> Chat with AI
"""

import asyncio
import os
import shutil
from pathlib import Path
from playwright.async_api import async_playwright

OUTPUT_DIR = Path("docs/images")
GIF_PATH = OUTPUT_DIR / "demo.gif"
FRONTEND_URL = "http://localhost:3000"
SAMPLE_DOC = Path("sample_docs/product_guide.txt").absolute()


async def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    frames_dir = OUTPUT_DIR / "frames"
    if frames_dir.exists():
        shutil.rmtree(frames_dir)
    frames_dir.mkdir()
    
    print("=" * 60)
    print("DEMO: KB + Document Upload + AI Chat")
    print("=" * 60)
    
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=False, slow_mo=30)
        ctx = await browser.new_context(viewport={"width": 1280, "height": 720})
        page = await ctx.new_page()
        
        fc = [0]
        async def snap(d="", n=1):
            for i in range(n):
                fc[0] += 1
                await page.screenshot(path=str(frames_dir / f"f_{fc[0]:03d}.png"))
                if i == 0 and d: print(f"  [{fc[0]:2d}] {d}")
                await asyncio.sleep(0.06)
        
        try:
            # 1. LANDING
            print("\n[1] Landing")
            await page.goto(FRONTEND_URL, timeout=20000)
            await asyncio.sleep(1)
            await snap("Info Naut", 3)
            
            # 2. DASHBOARD
            print("[2] Dashboard")
            await page.goto(f"{FRONTEND_URL}/dashboard", timeout=20000)
            await asyncio.sleep(1)
            await snap("Dashboard", 3)
            
            # 3. CREATE KB
            print("[3] Create Knowledge Base")
            await page.goto(f"{FRONTEND_URL}/dashboard/knowledge/new", timeout=20000)
            await asyncio.sleep(1)
            await snap("New KB", 2)
            
            await page.locator('input').first.fill("Product Docs")
            await asyncio.sleep(0.3)
            await snap("Name", 2)
            
            textarea = page.locator('textarea').first
            if await textarea.count() > 0:
                await textarea.fill("Product documentation")
            await snap("Description", 2)
            
            await page.click('button:has-text("Create")')
            await asyncio.sleep(2)
            await snap("KB Created!", 3)
            
            # Get KB ID from URL after redirect
            await asyncio.sleep(1)
            current_url = page.url
            print(f"    Current URL: {current_url}")
            
            # 4. UPLOAD DOCUMENT VIA API
            print("[4] Upload Document via API")
            
            # Get KB ID from the URL
            kb_id = current_url.split("/")[-1] if "/knowledge/" in current_url else "14"
            print(f"    KB ID: {kb_id}")
            
            # Upload via API
            import requests
            with open(SAMPLE_DOC, 'rb') as f:
                files = {'files': (SAMPLE_DOC.name, f, 'text/plain')}
                upload_resp = requests.post(
                    f"http://localhost:8000/api/knowledge-bases/{kb_id}/documents/upload",
                    files=files
                )
                print(f"    Upload response: {upload_resp.status_code}")
                if upload_resp.status_code == 200:
                    upload_data = upload_resp.json()
                    print(f"    Upload data: {upload_data}")
                    
                    # Process the document
                    if upload_data and not upload_data[0].get('skip_processing'):
                        process_resp = requests.post(
                            f"http://localhost:8000/api/knowledge-bases/{kb_id}/documents/process",
                            json=upload_data
                        )
                        print(f"    Process response: {process_resp.status_code}")
                        print("    Waiting for processing...")
                        await asyncio.sleep(5)
            
            # Show KB detail page
            await page.goto(f"{FRONTEND_URL}/dashboard/knowledge/{kb_id}", timeout=20000)
            await asyncio.sleep(2)
            await snap("KB with Document", 3)
            
            # Go to KB list
            await page.goto(f"{FRONTEND_URL}/dashboard/knowledge", timeout=20000)
            await asyncio.sleep(1.5)
            await snap("KB List", 3)
            
            # 5. CREATE CHAT VIA API
            print("[5] Create Chat via API")
            
            # Create chat via API
            import requests
            chat_resp = requests.post(
                "http://localhost:8000/api/chats/",
                json={"title": "Product Questions", "knowledge_base_ids": [int(kb_id)]}
            )
            print(f"    Chat creation: {chat_resp.status_code}")
            
            if chat_resp.status_code == 200:
                chat_data = chat_resp.json()
                chat_id = chat_data.get('id')
                print(f"    Chat ID: {chat_id}")
                
                # Navigate to chat page
                await page.goto(f"{FRONTEND_URL}/dashboard/chat/{chat_id}", timeout=20000)
                await asyncio.sleep(3)
                await snap("Chat Page", 3)
            else:
                print(f"    Chat error: {chat_resp.text}")
                await page.goto(f"{FRONTEND_URL}/dashboard/chat", timeout=20000)
                await asyncio.sleep(2)
                await snap("Chat List", 3)
            
            # 6. CHAT WITH AI
            print("[6] Chat with AI")
            await asyncio.sleep(2)
            await snap("Chat Interface", 3)
            
            msg = page.locator('input[placeholder="Type your message..."]')
            if await msg.count() > 0:
                # Q1
                q1 = "What are the key features of Info Naut?"
                print(f"    Q1: {q1}")
                await msg.fill(q1)
                await asyncio.sleep(0.5)
                await snap("Question 1", 3)
                
                await page.locator('button[type="submit"]').click()
                await snap("Sending", 2)
                
                print("    Waiting for AI response...")
                await asyncio.sleep(15)
                await snap("AI Response!", 6)
                
                await page.evaluate("window.scrollTo(0, 9999)")
                await asyncio.sleep(1)
                await snap("Full Response", 4)
                
                # Q2
                print("    Q2: Follow-up")
                await asyncio.sleep(2)
                msg2 = page.locator('input[placeholder="Type your message..."]')
                await msg2.fill("How does document processing work?")
                await asyncio.sleep(0.5)
                await snap("Question 2", 3)
                
                await page.locator('button[type="submit"]').click()
                print("    Waiting...")
                await asyncio.sleep(12)
                await snap("Response 2!", 5)
            
            # 7. API KEYS
            print("[7] API Keys")
            await page.goto(f"{FRONTEND_URL}/dashboard/api-keys", timeout=20000)
            await asyncio.sleep(1.5)
            await snap("API Keys", 3)
            
            # 8. DONE
            print("[8] Complete")
            await page.goto(f"{FRONTEND_URL}/dashboard", timeout=20000)
            await asyncio.sleep(1)
            await snap("Done!", 3)
            
        except Exception as e:
            print(f"Error: {e}")
            await page.screenshot(path="debug_error.png")
            await snap("Error", 2)
        
        await browser.close()
        
        print(f"\nFrames: {fc[0]}")
        
        # GIF
        print("\nCreating GIF...")
        from PIL import Image
        files = sorted(frames_dir.glob("f_*.png"))
        if files:
            frames = []
            for f in files:
                img = Image.open(f).resize((800, 450), Image.LANCZOS)
                img = img.convert('P', palette=Image.ADAPTIVE, colors=256)
                frames.append(img)
            
            frames[0].save(GIF_PATH, save_all=True, append_images=frames[1:],
                          duration=250, loop=0, optimize=True)
            
            sz = os.path.getsize(GIF_PATH) / (1024*1024)
            print(f"GIF: {sz:.2f} MB ({len(frames)} frames)")
            
            if sz > 5:
                opt = [frames[i] for i in range(0, len(frames), 2)]
                opt[0].save(GIF_PATH, save_all=True, append_images=opt[1:],
                           duration=500, loop=0, optimize=True)
            
            shutil.rmtree(frames_dir)
            print("\nSUCCESS!")


if __name__ == "__main__":
    asyncio.run(main())