Merge pull request #17 from alphagov/ACW-43/opensearch-integration

OpenSearch Integration

Author: daymohlar

app.py

@@ -1,22 +1,17 @@
 import asyncio
 import logging
 import os
-import time
 
 from asgiref.wsgi import WsgiToAsgi
 from dotenv import load_dotenv
 from flask import Flask, jsonify, render_template, request
 from werkzeug.exceptions import BadRequest
 
 from src.utils import (
-    background_run_extraction,
-    get_active_job_status,
-    get_job_id_for_path,
     read_job_status,
     resume_interrupted_jobs,
-    update_job_status,
+    start_extraction_job,
 )
-from src.visualiser_graph_generator import generate_output_path
 from src.visualiser_graph_loader import (
     extract_path_parts,
     load_json_file,
@@ -75,54 +70,21 @@ async def extract_quotes():
         """
         Endpoint that runs the Cytoscape graph generation logic based on graph.json.
         """
-        try:
-            source_path = request.args.get("source_path")
-            if not source_path:
-                return jsonify({"error": "Missing 'source_path' query parameter"}), 400
-
-            input_path, output_path = generate_output_path(source_path)
-            job_id = get_job_id_for_path(source_path)
-
-            active_status = get_active_job_status(job_id)
-            if active_status:
-                logger.info(
-                    f"Duplicate request for {source_path}. Job {job_id} is already in progress."
-                )
-                return jsonify(
-                    {
-                        "job_id": job_id,
-                        "status": "already_running",
-                        "message": (
-                            f"A graph generation job is already in progress for {source_path}"
-                        ),
-                        "output_path": output_path,
-                    }
-                ), 202
-
-            initial_status = {
-                "job_id": job_id,
-                "status": "pending",
-                "source_path": source_path,
-                "created_at": time.time(),
-            }
-            update_job_status(job_id, initial_status)
-
-            asyncio.create_task(
-                background_run_extraction(job_id, input_path, output_path, initial_status)
-            )
-
-            return jsonify(
-                {
-                    "job_id": job_id,
-                    "status": "accepted",
-                    "message": f"Graph generation started in background for {source_path}",
-                    "output_path": output_path,
-                }
-            ), 202
+        source_path = request.args.get("source_path") or ""
+        data, status_code = await start_extraction_job(source_path, extractor_type="s3")
+        return jsonify(data), status_code
 
-        except Exception as e:
-            app.logger.error(f"Error starting background task: {str(e)}")
-            return jsonify({"error": str(e)}), 500
+    @app.route("/extract-os", methods=["GET"])
+    async def extract_quotes_os():
+        """
+        Endpoint that runs the extraction using OpenSearch.
+        """
+        source_path = request.args.get("source_path") or ""
+        perform_indexing = request.args.get("index", "false").lower() == "true"
+        data, status_code = await start_extraction_job(
+            source_path, extractor_type="opensearch", perform_indexing=perform_indexing
+        )
+        return jsonify(data), status_code
 
     @app.route("/status/<job_id>", methods=["GET"])
     def get_status(job_id):

src/content_extractor/base.py

@@ -57,6 +57,10 @@ class BaseExtractorConfig:
     )
     chunk_max_chars: int = 6000
     secret_id: Optional[str] = None
+    aws_profile: Optional[str] = None
+    aws_access_key_id: Optional[str] = None
+    aws_secret_access_key: Optional[str] = None
+    aws_session_token: Optional[str] = None
 
 
 # --- Base Extractor Class ---
@@ -65,11 +69,23 @@ class BaseExtractorConfig:
 class BaseQuoteExtractor:
     def __init__(self, config: BaseExtractorConfig):
         self.config = config
-        self.s3_client = boto3.client("s3", region_name=self.config.region)
+
+        # Initialize AWS Session
+        self.session = boto3.Session(
+            profile_name=self.config.aws_profile,
+            aws_access_key_id=self.config.aws_access_key_id,
+            aws_secret_access_key=self.config.aws_secret_access_key,
+            aws_session_token=self.config.aws_session_token,
+            region_name=self.config.region,
+        )
+
+        self.s3_client = self.session.client("s3")
+        self.url_map: Dict[str, str] = {}
 
         # Initialize Bedrock Agent
+        bedrock_client = self.session.client("bedrock-runtime")
         model = BedrockConverseModel(
-            self.config.model_id, provider=BedrockProvider(region_name=self.config.region)
+            self.config.model_id, provider=BedrockProvider(bedrock_client=bedrock_client)
         )
         self.agent = Agent(
             model,
@@ -95,9 +111,43 @@ def __init__(self, config: BaseExtractorConfig):
             ),
         )
 
+    def _fetch_url_map(self, s3_uris: List[str]):
+        """
+        Attempts to fetch sources.json files from the directories of the input files.
+        Deduplicates potential sources.json locations and merges their mappings.
+        """
+        import logging
+
+        logger = logging.getLogger(__name__)
+
+        if not s3_uris:
+            return
+
+        sources_locations = set()
+        for uri in s3_uris:
+            if uri in self.url_map:
+                continue
+
+            if "/input/" in uri:
+                sources_uri = uri.split("/input/")[0] + "/input/sources.json"
+            else:
+                sources_uri = "/".join(uri.split("/")[:-1]) + "/sources.json"
+            sources_locations.add(sources_uri)
+
+        for sources_uri in sources_locations:
+            logger.info(f"Attempting to fetch sources map from {sources_uri}...")
+            content: Optional[str] = self.fetch_s3_content(sources_uri)
+            if content:
+                try:
+                    new_map = json.loads(content)
+                    self.url_map.update(new_map)
+                    logger.info(f"Successfully loaded {len(new_map)} mappings from {sources_uri}.")
+                except Exception as e:
+                    logger.error(f"Failed to parse {sources_uri}: {e}")
+
     def get_aws_secret(self, secret_id: str) -> dict:
         """Fetches and parses a JSON secret from AWS Secrets Manager."""
-        client = boto3.client("secretsmanager", region_name=self.config.region)
+        client = self.session.client("secretsmanager")
         try:
             response = client.get_secret_value(SecretId=secret_id)
             if "SecretString" in response: