import gradio as gr import asyncio import json import threading import time import sys import os import logging from typing import Optional, Dict, Any from datetime import datetime # Import from the published package from lf_algorithm import FrameworkAgent from lf_algorithm.utils import write_lineage_log # Configure logging for the demo server logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S' ) class SQLLineageFrontend: def __init__(self): self.agent_framework = None self.current_results = None self.current_agent_name = None self.log_thread = None self.should_stop_logging = False self.logger = logging.getLogger(__name__) def get_visualize_link(self) -> str: """Generate JSONCrack visualization interface for aggregation data""" if self.current_results is None: return """
📊 Visualization Ready
After you run analysis and succeed, you need to got to the following link:
🔗 Open editor for simple check and paste the results there
""" try: # Get the aggregation data - now it's directly the current_results aggregation_data = self.current_results # Handle different result types if isinstance(aggregation_data, str): try: # Try to parse as JSON first parsed_data = json.loads(aggregation_data) data_to_encode = parsed_data except json.JSONDecodeError: # If it's not valid JSON, wrap it in a dict data_to_encode = {"aggregation_output": aggregation_data} elif hasattr(aggregation_data, 'to_dict'): # Handle AgentResult objects data_to_encode = aggregation_data.to_dict() elif isinstance(aggregation_data, dict): data_to_encode = aggregation_data else: # Fallback for other object types data_to_encode = {"aggregation_output": str(aggregation_data)} # Format JSON for display formatted_json = json.dumps(data_to_encode, indent=2) return f"""
✅ Analysis Complete! Ready for Visualization
📋 Steps to visualize your results:
1. Click "Open JSONCrack Editor" below
2. Click "Copy JSON" button or click the JSON data below to select all
3. Paste it into the JSONCrack editor
🔗 Open JSONCrack Editor

📄 Analysis Results (JSON)
""" except Exception as e: return f"
❌ Error generating visualization data: {str(e)}
" def get_logs_html(self) -> str: """Generate HTML for live logs display""" if self.current_agent_name is None: return "
No agent initialized yet
" return f"""
📝 Logging Status for Agent: {self.current_agent_name}
Standard Python Logging Active
• All logs are being captured by the application's logging system
• Check your console/terminal for real-time log output
• Logs include detailed information about agent execution
• Structured logging with timestamps and log levels

📋 Log Types Available:
INFO - General information and progress
DEBUG - Detailed debugging information
WARNING - Warning messages
ERROR - Error messages

🔍 What You'll See:
• Agent initialization and configuration
• MCP tool interactions and responses
• Analysis progress and completion status
• Any errors or warnings during execution
""" def test_log_writing(self): """Test function to write a sample log entry""" if self.current_agent_name: try: write_lineage_log(self.current_agent_name, "test", "Test log entry from frontend") self.logger.info(f"Test log written successfully for agent: {self.current_agent_name}") return f"✅ Test log written successfully for agent: {self.current_agent_name}! Check your console output." except Exception as e: self.logger.error(f"Failed to write test log: {e}") return f"❌ Failed to write test log: {e}" else: return "⚠️ Please initialize an agent first by running an analysis" def get_results_info(self) -> str: """Get information about the current results""" if self.current_results is None: return "No results available yet" if isinstance(self.current_results, dict) and "error" in self.current_results: return f"Error in results: {self.current_results['error']}" if hasattr(self.current_results, 'to_dict'): # AgentResult object result_dict = self.current_results.to_dict() inputs_count = len(result_dict.get('inputs', [])) outputs_count = len(result_dict.get('outputs', [])) return f"✅ Structured results with {inputs_count} input(s) and {outputs_count} output(s)" if isinstance(self.current_results, dict): return f"✅ Dictionary results with {len(self.current_results)} keys" return f"✅ Results type: {type(self.current_results).__name__}" async def run_analysis(self, agent_name: str, model_name: str, query: str): """Run SQL lineage analysis""" try: # Validate input if not query or not query.strip(): return "❌ Error: Query cannot be empty. Please provide a valid query for analysis." self.logger.info(f"Starting analysis with agent: {agent_name}, model: {model_name}") # Initialize the agent framework with simplified constructor self.agent_framework = FrameworkAgent( agent_name=agent_name, model_name=model_name, source_code=query.strip() ) self.current_agent_name = agent_name self.logger.info(f"Agent framework initialized. Running analysis...") # Run the analysis using the structured results method results = await self.agent_framework.run_agent() self.current_results = results # Check if we got an error response if isinstance(results, dict) and "error" in results: self.logger.error(f"Analysis failed: {results['error']}") return f"❌ Analysis failed: {results['error']}" self.logger.info(f"Analysis completed successfully for agent: {agent_name}") return f"""✅ Analysis completed successfully! Results are now available in the visualization section. Click 'Open JSONCrack Editor' to visualize your data lineage. If you want to set up your own local development environment or deploy this in production, please refer to the GitHub repository mentioned above.""" except ValueError as ve: self.logger.error(f"Validation error: {ve}") return f"❌ Validation error: {str(ve)}" except Exception as e: self.logger.error(f"Error running analysis: {e}") return f"❌ Error running analysis: {str(e)}" def run_analysis_sync(self, agent_name: str, model_name: str, query: str): """Synchronous wrapper for run_analysis""" return asyncio.run(self.run_analysis(agent_name, model_name, query)) def create_ui(self): """Create the Gradio interface""" with gr.Blocks(title="SQL Lineage Analysis", fill_width=True) as ui: gr.Markdown('
🔍 Demo Lineagentic-Flow
') gr.Markdown('
Analyze data lineage with AI-powered agents
') gr.Markdown('
Check out agent types for supporting script types
') gr.Markdown('
For local and production runs, check out the repo: 🔗 https://github.com/lineagentic/lineagentic-flow
') with gr.Row(): # Left column - Configuration and Query with gr.Column(scale=1): gr.Markdown("### 1. Agent Configuration") agent_dropdown = gr.Dropdown( label="Agent Type", choices=[ "sql-lineage-agent", "python-lineage-agent", "airflow-lineage-agent", "java-lineage-agent", "spark-lineage-agent" ], value="sql-lineage-agent" ) model_dropdown = gr.Dropdown( label="Model", choices=[ "gpt-4o-mini", "gpt-4o", "deepseek-coder", "deepseek-chat", "gemini-pro" ], value="gpt-4o-mini" ) gr.Markdown("### 2. Query for Lineage Analysis") query_input = gr.Textbox( label="Query", placeholder="Enter your SQL query here...", lines=9, max_lines=15 ) analyze_button = gr.Button("🚀 Run Analysis", variant="primary", size="lg") status_output = gr.Textbox(label="Status", interactive=False) # Right column - Visualization and Logs with gr.Column(scale=1): gr.Markdown("### 3. Results Information") results_info = gr.Textbox( label="Results Status", value=self.get_results_info(), interactive=False ) gr.Markdown("### 4. Visualize Results") gr.Markdown("📊 After successful analysis, visualize your results in demo editor") visualize_html = gr.HTML(self.get_visualize_link()) gr.Markdown("### 5. Logging Information") logs_html = gr.HTML(self.get_logs_html()) test_log_button = gr.Button("Test Log Writing", variant="secondary", size="sm") # Auto-refresh logs every 5 seconds refresh_logs = gr.Button("🔄 Refresh Logs", variant="secondary", size="sm") refresh_results = gr.Button("🔄 Refresh Results Info", variant="secondary", size="sm") # Event handlers def run_analysis_and_update(agent_name, model_name, query): """Run analysis and update visualization""" # Run the analysis status_result = self.run_analysis_sync(agent_name, model_name, query) # Update visualization, logs, and results info viz_html = self.get_visualize_link() logs_html = self.get_logs_html() results_info = self.get_results_info() return status_result, results_info, viz_html, logs_html analyze_button.click( fn=run_analysis_and_update, inputs=[agent_dropdown, model_dropdown, query_input], outputs=[status_output, results_info, visualize_html, logs_html] ) test_log_button.click( fn=self.test_log_writing, inputs=[], outputs=[status_output] ) refresh_logs.click( fn=self.get_logs_html, inputs=[], outputs=[logs_html] ) refresh_results.click( fn=self.get_results_info, inputs=[], outputs=[results_info] ) return ui def run(self): """Launch the Gradio interface""" ui = self.create_ui() ui.launch(share=False, inbrowser=True) if __name__ == "__main__": frontend = SQLLineageFrontend() frontend.run()