feat(pdf): add cover page and builder class (#19)

Author: NeurArk

README.md

@@ -40,6 +40,27 @@ docker run -p 7860:7860 neurark/mcp-data-assistant
 
 ▶ **Demo** — open [`sample_docs/report-demo.pdf`](sample_docs/report-demo.pdf)
 
+### Advanced PDF reports
+The PDF tool now supports a cover page with optional logo and summary box as
+well as styled charts. A minimal example:
+
+```json
+{
+  "title": "Sales Overview",
+  "summary": "Highlights for 2024",
+  "cover": {"logo_path": "assets/logo.png"},
+  "sections": [
+    {"title": "Intro", "type": "paragraph", "text": "Quarterly results"},
+    {
+      "title": "Revenue",
+      "type": "chart",
+      "chart_spec": {"chart_type": "bar", "labels": ["Q1", "Q2"], "values": [1, 2], "color": "#ff8800"}
+    }
+  ]
+}
+```
+Run `create_pdf` with this JSON to produce a multi-page report with a styled chart.
+
 ## How it works
 The app launches Gradio with `mcp_server=True`.  
 The LLM discovers three tools via the MCP schema and chains them as

sample_docs/report-demo-v2.pdf

@@ -19,6 +19,10 @@
       },
       "additionalProperties": false
     },
+    "summary": {
+      "type": "string",
+      "description": "Short executive summary displayed on the cover page"
+    },
     "insights": {
       "type": "array",
       "description": "List of insight paragraphs to highlight on the first page",
@@ -35,14 +39,13 @@
           "text": {"type": "string"},
           "data": {"type": ["array", "object"]},
           "chart_spec": {
-            "type": "object",
-            "properties": {
-              "chart_type": {"type": "string", "enum": ["bar", "pie", "line"]},
-              "labels": {"type": "array", "items": {"type": "string"}},
-              "values": {"type": "array", "items": {"type": "number"}}
-            },
-            "required": ["chart_type", "labels", "values"],
-            "additionalProperties": false
+            "oneOf": [
+              {"$ref": "#/definitions/chartSpec"},
+              {
+                "type": "array",
+                "items": {"$ref": "#/definitions/chartSpec"}
+              }
+            ]
           }
         },
         "required": ["title", "type"],
@@ -51,5 +54,20 @@
     }
   },
   "required": ["title", "sections"],
-  "additionalProperties": false
+  "additionalProperties": false,
+  "definitions": {
+    "chartSpec": {
+      "type": "object",
+      "properties": {
+        "chart_type": {"type": "string", "enum": ["bar", "pie", "line"]},
+        "labels": {"type": "array", "items": {"type": ["string", "number"]}},
+        "values": {"type": "array", "items": {"type": "number"}},
+        "color": {"type": "string", "description": "Hex color"},
+        "width": {"type": "number", "description": "Figure width in inches"},
+        "height": {"type": "number", "description": "Figure height in inches"}
+      },
+      "required": ["chart_type", "labels", "values"],
+      "additionalProperties": false
+    }
+  }
 }

static/pdf_schema.json

@@ -4,15 +4,15 @@
 from agent.ollama_integration import (
     check_ollama_available,
     get_ollama_model_name,
-    create_ollama_model
+    create_ollama_model,
 )
 from agent import answer
 import httpx
 
 # Suppress the coroutine warning for tests
-warnings.filterwarnings("ignore",
-                       message="coroutine '.*' was never awaited",
-                       category=RuntimeWarning)
+warnings.filterwarnings(
+    "ignore", message="coroutine '.*' was never awaited", category=RuntimeWarning
+)
 
 
 def test_check_ollama_available():
@@ -58,7 +58,7 @@ def test_ollama_provider():
         mock_result.final_output = "The answer to 2+2 is 4"
 
         # Patch the asyncio run function to avoid actual API calls
-        with patch('agent.assistant.asyncio.run', return_value=mock_result):
+        with patch("agent.assistant.asyncio.run", return_value=mock_result):
             # This should now run without errors since the API call is mocked
             response, _ = answer("What is 2+2?", provider="ollama")
 
@@ -97,6 +97,7 @@ def test_ollama_provider():
 
     except Exception as e:
         import traceback
+
         print(f"❌ Error testing Ollama provider: {str(e)}")
         print(traceback.format_exc())
         raise
@@ -113,6 +114,7 @@ def test_ollama_tool_knowledge():
         # Create a session to maintain context
         from agent.session_manager import session_manager
         from agents.mcp import MCPServerSse
+
         session_id = session_manager.create_session()
         print(f"\nCreated test session: {session_id}")
 
@@ -188,32 +190,45 @@ async def connect_mcp_server():
         found_terms = [term for term in expected_terms if term in response.lower()]
 
         print(f"Found terms: {found_terms}")
-        assert found_terms, f"Expected tool terms ({', '.join(expected_terms)}) not found in response"
+        assert found_terms, (
+            f"Expected tool terms ({', '.join(expected_terms)}) not found in response"
+        )
 
         # Check for specific error messages
-        assert "[] is too short - 'messages'" not in response, "Error: Empty messages array sent to Ollama API"
+        assert "[] is too short - 'messages'" not in response, (
+            "Error: Empty messages array sent to Ollama API"
+        )
 
         # Step 3: Test conversation history and follow-up
         print("\n3️⃣ Testing follow-up question...")
         followup_query = "Can you list the tools again and explain what each one does?"
-        followup_response, _ = answer(followup_query, provider="ollama", session_id=session_id)
+        followup_response, _ = answer(
+            followup_query, provider="ollama", session_id=session_id
+        )
         print(f"Follow-up response: {followup_response[:500]}...")
 
         # Verify the response has relevant terms
         followup_terms = ["csv", "sql", "pdf", "database", "file"]
-        found_followup_terms = [term for term in followup_terms if term in followup_response.lower()]
-        assert found_followup_terms, f"Follow-up response doesn't contain expected terms"
+        found_followup_terms = [
+            term for term in followup_terms if term in followup_response.lower()
+        ]
+        assert found_followup_terms, "Follow-up response doesn't contain expected terms"
 
         # Final check: Conversation history maintained
         history = session_manager.get_messages(session_id)
         print(f"\n✅ Session maintained context through {len(history)} messages")
-        assert len(history) >= 4, "Expected at least 4 messages in conversation history (2 queries + 2 responses)"
+        assert len(history) >= 4, (
+            "Expected at least 4 messages in conversation history (2 queries + 2 responses)"
+        )
 
         print("\n✅ Direct Ollama integration test PASSED.")
-        print(f"   Found terms in responses: {', '.join(found_terms + found_followup_terms)}")
+        print(
+            f"   Found terms in responses: {', '.join(found_terms + found_followup_terms)}"
+        )
 
     except Exception as e:
         import traceback
+
         print(f"\n❌ Direct Ollama integration test FAILED: {str(e)}")
         print(traceback.format_exc())
         raise
@@ -258,13 +273,13 @@ def test_provider_fallback():
     # Test fallback if Ollama is unavailable
     if not check_ollama_available():
         response, result = answer("Test", provider="ollama")
-        assert (
-            "⚠️ Ollama not available" in response
-        ), "Expected unavailable message for Ollama provider"
+        assert "⚠️ Ollama not available" in response, (
+            "Expected unavailable message for Ollama provider"
+        )
 
     # Test fallback if OpenAI key is not set
     if os.getenv("OPENAI_API_KEY") is None:
         response, result = answer("Test", provider="openai")
-        assert (
-            "⚠️ OPENAI_API_KEY not set" in response
-        ), "Expected API key message for OpenAI provider"
+        assert "⚠️ OPENAI_API_KEY not set" in response, (
+            "Expected API key message for OpenAI provider"
+        )

tests/test_ollama_integration.py

@@ -115,7 +115,7 @@ def test_pdf_with_edge_cases(tmp_path):
     """Test PDF creation with edge cases."""
     # Long text
     long_text = (
-        "This is a very long text that should be wrapped properly " "in the PDF table "
+        "This is a very long text that should be wrapped properly in the PDF table "
     ) * 5
     data = {
         "title": "Edge Case Test",
@@ -238,6 +238,56 @@ def _count_images(pdf_path: Path) -> int:
         return f.read().count(b"/Subtype /Image")
 
 
+def test_pdf_with_cover_and_summary(tmp_path):
+    data = {
+        "title": "Cover Report",
+        "summary": "Quick overview",
+        "cover": {"logo_path": "assets/logo.png"},
+        "sections": [{"title": "Intro", "type": "paragraph", "text": "Hello"}],
+    }
+    pdf_path = Path(create_pdf(data, out_path=tmp_path / "cover.pdf"))
+    assert pdf_path.exists()
+    from PyPDF2 import PdfReader
+
+    reader = PdfReader(str(pdf_path))
+    assert len(reader.pages) >= 2
+    assert "Quick overview" in reader.pages[0].extract_text()
+
+
+def test_multiple_chart_specs(tmp_path):
+    data = {
+        "title": "Charts",
+        "sections": [
+            {
+                "title": "Multi",
+                "type": "chart",
+                "chart_spec": [
+                    {
+                        "chart_type": "bar",
+                        "labels": ["A", "B"],
+                        "values": [1, 2],
+                        "color": "#ff0000",
+                    },
+                    {"chart_type": "line", "labels": [1, 2], "values": [3, 4]},
+                ],
+            }
+        ],
+    }
+    pdf_path = Path(create_pdf(data, out_path=tmp_path / "multi.pdf"))
+    assert pdf_path.exists()
+    assert _count_images(pdf_path) >= 3
+
+
+def test_builder_class(tmp_path):
+    from tools.pdf_tool import PdfReportBuilder
+
+    with PdfReportBuilder(tmp_path / "builder.pdf") as builder:
+        builder.add_cover("Title")
+        builder.add_section({"title": "P", "type": "paragraph", "text": "Hi"})
+        path = builder.save()
+    assert Path(path).exists()
+
+
 def test_pdf_with_sections_and_charts(tmp_path):
     """Generate a PDF using the new schema with multiple chart types."""
     data = {