feat: integrated LLM Assistant for PDF reports (#12)

* feat: integrated LLM Assistant for PDF reports

- Added agent with OpenAI integration to handle natural language requests
- Enhanced PDF generation tool to handle nested data
- Added special handling for data from the agent
- Fixed several display issues with PDF reports
- Extended test suite for better error handling
- Improved demo scripts for CLI testing

* fix: fix linting issues with flake8

- Fixed missing whitespace around arithmetic operators in app.py and pdf_tool.py
- Fixed line break positioning for binary operators in tests and pdf_tool.py
- Updated GitHub CI configuration to ignore conflicting rules W503 and W504

Author: NeurArk

.github/workflows/ci.yml

@@ -14,7 +14,7 @@ jobs:
         with:
           python-version: "3.12"
       - run: pip install flake8
-      - run: flake8 app.py tools/ tests/ --max-line-length 100 --exclude=".venv/" --ignore=E302,W293,W291,E128,W292,F401,F841,E305
+      - run: flake8 app.py tools/ tests/ --max-line-length 100 --exclude=".venv/" --ignore=E302,W293,W291,E128,W292,F401,F841,E305,W503,W504
 
   tests:
     runs-on: ubuntu-latest

agent/__init__.py

@@ -0,0 +1,7 @@
+"""
+Agent integration module for the MCP Data Assistant.
+"""
+
+from agent.assistant import answer
+
+__all__ = ["answer"]

agent/assistant.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+import os
+import asyncio
+from agents import Agent, Runner
+from agents.mcp import MCPServerSse
+
+# Gradio 5.29 SSE endpoint
+MCP_SSE_URL = os.getenv(
+    "MCP_SSE_URL",
+    "http://127.0.0.1:7860/gradio_api/mcp/sse",
+)
+
+# Create MCP server instance but don't connect yet
+mcp_server = MCPServerSse(
+    params={"url": MCP_SSE_URL},
+    cache_tools_list=True,
+)
+
+# Initialize agent
+agent = Agent(
+    name="NeurArk Data Assistant",
+    instructions=(
+        "You are a data assistant that can analyze tabular data and create PDFs.\n"
+        "You can work with SQL databases, CSV files, and generate PDF reports.\n"
+        "Common workflows include:\n"
+        "- Query data from database then generate PDF report with results\n"
+        "- Analyze CSV files and create summary reports\n"
+        "- Generate custom reports based on user specifications\n"
+        "You should auto-discover available tools via the MCP server connection.\n\n"
+        "When working with databases:\n"
+        "- First, discover what tables are available in the database\n"
+        "- If the user mentions a table that doesn't exist, look for alternatives\n"
+        "- Explore the structure of the tables to understand columns\n"
+        "- Execute appropriate queries based on what you discovered\n\n"
+        "When generating PDF reports:\n"
+        "- The 'data_json' parameter should be a JSON string with data to include\n"
+        "- Always include the generated PDF file path in your response\n"
+        "- Example format: {\"title\": \"Report Title\", \"data\": \"Your Data\"}\n"
+    ),
+    model="gpt-4.1-mini",
+    mcp_servers=[mcp_server],
+)
+
+
+async def _run_agent(prompt: str) -> str:
+    """Run the agent asynchronously with proper server connection handling."""
+    # Connect to MCP server before running the agent
+    async with mcp_server:
+        # Execute the agent with the prompt
+        result = await Runner.run(starting_agent=agent, input=prompt)
+        return result.final_output  # String with PDF path or response
+
+
+def answer(prompt: str) -> str:
+    """Synchronous wrapper for running the agent."""
+    if not os.getenv("OPENAI_API_KEY"):
+        return "⚠️ OPENAI_API_KEY not set."
+
+    try:
+        # Run the async function in a synchronous context
+        return asyncio.run(_run_agent(prompt))
+    except Exception as e:
+        import traceback
+        error_trace = traceback.format_exc()
+        print(f"Agent error: {str(e)}")
+        print(f"Error trace: {error_trace}")
+        return f"Error: {str(e)}\nTrace: {error_trace}"

app.py

@@ -1,21 +1,33 @@
+# existing imports
 import gradio as gr
+import json
+import pathlib
+import threading
+import time
+import requests
 from tools.sql_tool import run_sql
 from tools.csv_tool import summarise_csv
 from tools.pdf_tool import create_pdf
+# assistant
+from agent import answer
+
 
 def server_status() -> str:
     """
     A dummy function to show the server is alive.
-    
+
     Returns:
         str: Status message confirming the server is running
     """
+    # Function attribute to hide from MCP
+    server_status._hide_from_mcp = True
     return "✅ MCP Data Assistant server is running."
 
-with gr.Blocks() as demo:
+
+with gr.Blocks() as tools_demo:
     gr.Markdown("# MCP Data Assistant")
-    gr.Markdown("This server will expose three tools (SQL query, CSV summary, PDF report).")
-    
+    gr.Markdown("This server will expose three tools (SQL, CSV summary, PDF report).")
+
     # Register MCP tools
     run_sql_interface = gr.Interface(
         fn=run_sql,
@@ -26,7 +38,7 @@ def server_status() -> str:
         examples=["SELECT 1 AS one"],
         api_name="sql"
     )
-    
+
     summarise_csv_interface = gr.Interface(
         fn=summarise_csv,
         inputs=gr.Textbox(label="CSV File Path"),
@@ -36,47 +48,181 @@ def server_status() -> str:
         examples=["sample_data/people.csv"],
         api_name="csv"
     )
-    
+
+    # Wrapper around create_pdf to ensure data parameter is properly processed
+    def create_pdf_wrapper(data_json, out_path=None, include_chart=True):
+        """
+        Generate a professional PDF report from provided data.
+
+        Creates a PDF document with the given data formatted as a table.
+        Optionally includes a bar chart visualization of numeric values.
+
+        Args:
+            data_json: JSON string or object containing the data to include
+            out_path: Optional custom path for the generated PDF file
+            include_chart: Whether to include a bar chart visualization
+
+        Returns:
+            Absolute path to the generated PDF file
+
+        Raises:
+            ValueError: If the data dictionary is empty
+        """
+        # Debug log (minimal)
+        print(f"PDF request received with type: {type(data_json)}")
+
+        # Parse JSON string to dict if needed
+        if isinstance(data_json, str):
+            try:
+                data = json.loads(data_json)
+            except Exception:
+                # Handle invalid JSON by creating an error dict
+                data = {
+                    "error": "Invalid JSON",
+                    "raw_input": (data_json[:200] + "..."
+                               if len(data_json) > 200 else data_json)
+                }
+        else:
+            # Use the data directly
+            data = data_json
+
+        try:
+            # Handle basic data type conversion
+            if isinstance(data, dict):
+                # Dictionary - use as is
+                pass
+            elif isinstance(data, list):
+                # Convert list to simple dictionary with indexed keys
+                items = {"item_" + str(i + 1): item for i, item in enumerate(data)}
+                data = items
+            else:
+                # Unsupported type - create error dict
+                data = {
+                    "error": "Unsupported data type",
+                    "received_type": str(type(data))
+                }
+
+            # Create the PDF
+            return create_pdf(data, out_path, include_chart)
+
+        except Exception as e:
+            # If PDF creation fails, create an error report
+            try:
+                error_data = {"error": f"Failed to create PDF: {str(e)}"}
+                return create_pdf(error_data, out_path, include_chart=False)
+            except Exception:
+                # Last resort if even the error PDF can't be created
+                return "Critical error creating PDF"
+
     create_pdf_interface = gr.Interface(
-        fn=create_pdf,
+        fn=create_pdf_wrapper,
         inputs=[
-            gr.JSON(label="Report Data"),
-            gr.Textbox(label="Output Path (optional)", 
-                    placeholder="Leave empty for default location"),
+            gr.Textbox(
+                label="Report Data (JSON)",
+                value='{"customer": "ACME", "total": 1000}'
+            ),
+            gr.Textbox(
+                label="Output Path (optional)",
+                placeholder="Leave empty for default location"
+            ),
             gr.Checkbox(label="Include Chart", value=True)
         ],
         outputs=gr.Textbox(label="Generated PDF Path"),
         title="PDF Report Generator",
         description="Create professional PDF reports with data and optional charts",
         examples=[
-            [{"customer": "ACME", "order_id": 12345, "total": 999, "items": 5}, None, True]
+            ['{"customer": "ACME", "total": 999}',
+             None, True]
         ],
         api_name="pdf"
     )
-    
+
     # Add simple UI components
     status_btn = gr.Button("Ping server")
     status_output = gr.Textbox()
-    status_btn.click(server_status, outputs=status_output)
+    # Hide from API and MCP
+    status_btn.click(server_status, outputs=status_output, api_name=False)
+
+
+# ---------- Assistant tab ----------
+assistant_chat = gr.ChatInterface(
+    fn=answer,
+    title="NeurArk Data Assistant",
+    examples=[
+        "Show me total sales for 2024 and create a PDF report"
+    ],
+    api_name=False,  # Hide from API and MCP
+    # Specify message type explicitly to avoid warning
+    chatbot=gr.Chatbot(type="messages"),
+    type="messages"
+)
+
+
+# ---------- Tabs UI -----------------
+demo = gr.TabbedInterface(
+    [tools_demo, assistant_chat],
+    ["Tools demo", "Assistant"],
+    title="NeurArk MCP Data Assistant",
+)
+
+
+# Create a function to save the schema with retry
+def save_schema_with_retry(retries=3, delay=0.5):
+    """Try to save the schema with retries in case the server isn't ready yet."""
+    for attempt in range(retries):
+        try:
+            pathlib.Path("static").mkdir(exist_ok=True)
+            # Use a fixed URL
+            schema_url = "http://127.0.0.1:7860/gradio_api/mcp/schema"
+            response = requests.get(schema_url, timeout=2)  # Short timeout
+            if response.status_code == 200:
+                schema = response.json()
+
+                # Keep only the tools we want to expose
+                filtered_schema = {
+                    k: v for k, v in schema.items()
+                    if k in ["sql", "csv", "pdf"]
+                }
+
+                with open("static/schema.json", "w") as f:
+                    f.write(json.dumps(filtered_schema, indent=2))
+                print("Schema saved to static/schema.json")
+
+                # For information, display available tools
+                tools_list = ', '.join(filtered_schema.keys())
+                print(f"MCP Tools available: {tools_list}")
+                return filtered_schema
+        except Exception as e:
+            if attempt < retries - 1:
+                print(
+                    f"Attempt {attempt + 1}/{retries} failed: {e}. "
+                    f"Retrying in {delay}s..."
+                )
+                time.sleep(delay)
+            else:
+                print(f"Failed to save schema after {retries} attempts: {e}")
+    return None
+
+
+# Function to save schema in background
+def background_save_schema():
+    # Wait for server to start
+    time.sleep(2.0)
+    # Try to save the schema
+    save_schema_with_retry(retries=3, delay=1.0)
+
 
 if __name__ == "__main__":
-    # Ensure the static folder exists
-    import os
-    import pathlib
-    import json
-    import requests
-    
+    # Create and start a thread to save the schema in the background
+    schema_thread = threading.Thread(target=background_save_schema)
+    schema_thread.daemon = True  # Thread will stop when main program stops
+    schema_thread.start()
+
+    print("Starting MCP server...")
+
     # Enable MCP server for LLM tools access
-    app = demo.launch(mcp_server=True, share=False, show_error=True)
-    
-    # After launch, fetch and save the schema
-    try:
-        pathlib.Path("static").mkdir(exist_ok=True)
-        schema_url = f"{app.local_url}gradio_api/mcp/schema"
-        response = requests.get(schema_url)
-        if response.status_code == 200:
-            with open("static/schema.json", "w") as f:
-                f.write(json.dumps(response.json(), indent=2))
-            print("Schema saved to static/schema.json")
-    except Exception as e:
-        print(f"Failed to save schema: {e}")
+    # In Gradio 5.29, launch the server in a blocking way (default)
+    demo.launch(mcp_server=True, share=False, show_error=True)
+
+    # This code will never be reached because launch() is blocking
+    print("Server stopped.")