Merge pull request #7 from alphagov/ACW-32/store-graph-to-s3

ACW-32 Add functionality to read and write graphs to S3

Author: daymohlar

app.py

@@ -2,7 +2,7 @@
 import logging
 from flask import Flask, request, jsonify, render_template
 from dotenv import load_dotenv
-from src.generate_graph import generate_graph, load_graph_viewmodel
+from src.generate_graph import generate_graph, load_graph_viewmodel, generate_output_path, summarize_path
 
 load_dotenv()
 
@@ -47,11 +47,15 @@ async def extract_quotes():
         Endpoint that runs the Cytoscape graph generation logic based on graph.json.
         """
         try:
-            logger.info('Starting graph generation process...')
-            graph_data = await generate_graph("graph.json")
+            input_path = request.args.get('input_path')
+            if not input_path:
+                return jsonify({"error": "Missing 'input_path' query parameter"}), 400
+
+            output_path = generate_output_path(input_path)
+            logger.info(f'Starting graph generation process for {summarize_path(input_path)}...')
+            graph_data = await generate_graph(input_path, output_path)
             logger.info('Graph generation completed successfully.')
-            
-            return jsonify(graph_data), 200
+            return jsonify({'status': 'running'}), 200
 
         except Exception as e:
             app.logger.error(f"Error generating graph: {str(e)}")

src/generate_graph.py

@@ -8,10 +8,12 @@
 from src.content_extractor.s3_sequential import S3QuoteExtractor
 from src.content_extractor.base import BaseExtractorConfig
 from src.content_extractor.highlighter import highlight_occurrence
+import fsspec
 from src.models.graph_models import (
     GraphInput, GraphOutput, Node, NodeData, Edge, EdgeData, Occurrence, Entity
 )
 
+
 logger = logging.getLogger(__name__)
 
 def slugify(text: str) -> str:
@@ -124,11 +126,12 @@ def build_node_structure(entities: List[Entity], entity_results: Dict[str, Any])
 async def generate_graph(input_data: Union[str, Dict[str, Any]], output_path: Optional[str] = None):
     """Main orchestration function. Can take a file path (str) or a dictionary."""
     if isinstance(input_data, str):
-        if not os.path.exists(input_data):
+        try:
+            with fsspec.open(input_data, "r") as f:
+                graph_data = json.load(f)
+        except FileNotFoundError:
             logger.error(f"Input file {input_data} not found.")
-            return
-        with open(input_data, "r") as f:
-            graph_data = json.load(f)
+            raise
     else:
         graph_data = input_data
 
@@ -149,10 +152,9 @@ async def generate_graph(input_data: Union[str, Dict[str, Any]], output_path: Op
     cy_json = cy_graph.model_dump(exclude_none=True)
 
     if output_path:
-        os.makedirs(os.path.dirname(output_path), exist_ok=True)
-        with open(output_path, "w") as f:
+        with fsspec.open(output_path, "w", auto_mkdir=True) as f:
             json.dump(cy_json, f, indent=4)
-        logger.info(f"Graph saved to {output_path}")
+        logger.info(f"Graph saved to {summarize_path(output_path)}")
 
     return cy_json
 
@@ -168,5 +170,32 @@ def load_graph_viewmodel(file_path: str) -> Dict[str, Any]:
     """Loads the graph viewmodel JSON for the frontend."""
     return load_json_file(file_path)
 
+
+def summarize_path(path: str) -> str:
+    """Extracts a concise representation (project/run) of a path for logging."""
+    match = re.search(r'([^/]+)/(run-\d+-\d+)', path)
+    if match:
+        return f"{match.group(1)}/{match.group(2)}"
+    return path.split('/')[-1]
+
+def generate_output_path(input_path: str) -> str:
+    """Generates the output path for the graph JSON file."""
+    
+    #TODO: make input from user be relative without the bucketname applied
+    match = re.search(r'(?P<base>s3://govuk-ai-accelerator-data-integration/)?(?P<domain_name>[^/]+)/(?P<run>run-\d+-\d+)', input_path)
+    
+    if match:
+        base = match.group('base')
+        domain_name = match.group('domain_name')
+        run_id = match.group('run')
+        
+        if base:
+            return f"{base}graph_tools/{domain_name}/{run_id}/graphNode.json"
+        else:
+            raise ValueError(f"Input path '{input_path}' does not contain a recognizable S3 path structure.")
+    
+    summary = summarize_path(input_path)
+    raise ValueError(f"Input path '{summary}' does not contain a recognizable project/run structure.")
+
 if __name__ == "__main__":
     asyncio.run(generate_graph("graph.json", "outputs/graphNode.json"))