modcomms/backend/app/services/knowledge_base_service.py

import logging
import uuid
from datetime import datetime, timezone
from typing import Optional

from app.models.database import async_session_factory
from app.repositories.knowledge_base_repository import KnowledgeBaseRepository
from app.services.storage_service import storage_service

logger = logging.getLogger(__name__)

# Distillation prompt templates per agent type
DISTILLATION_PROMPTS = {
    "legal": """You are a compliance documentation specialist. Below is raw reference material about legal compliance, advertising standards, financial promotions, and disclaimers relevant to Barclays marketing materials.

Your task is to distil this into a clear, structured specification document that an AI compliance agent can use to review marketing proofs. Organise the content into logical sections with clear headings. Focus on actionable rules, required disclaimers, prohibited content, and compliance requirements.

Remove any redundancy, marketing fluff, or content not relevant to compliance checking. Preserve all specific rules, thresholds, and requirements verbatim.

RAW REFERENCE MATERIAL:
{combined_markdown}

OUTPUT: A well-structured markdown specification document.""",

    "brand_barclays": """You are a brand guidelines specialist. Below is raw reference material about Barclays brand guidelines including logo usage, colour palettes, typography, design principles, and visual identity standards.

Your task is to distil this into a clear, structured specification document that an AI brand compliance agent can use to review marketing proofs. Organise the content into logical sections: logo rules, colour specifications (with exact hex/RGB values), typography rules, spacing/layout requirements, do's and don'ts.

Remove any redundancy or content not directly relevant to visual brand compliance checking. Preserve all specific measurements, colour values, and rules verbatim.

RAW REFERENCE MATERIAL:
{combined_markdown}

OUTPUT: A well-structured markdown specification document.""",

    "brand_barclaycard": """You are a brand guidelines specialist. Below is raw reference material about Barclaycard brand guidelines including logo usage, colour palettes, typography, design principles, and visual identity standards.

Your task is to distil this into a clear, structured specification document that an AI brand compliance agent can use to review marketing proofs. Organise the content into logical sections: logo rules, colour specifications (with exact hex/RGB values), typography rules, spacing/layout requirements, do's and don'ts.

Remove any redundancy or content not directly relevant to visual brand compliance checking. Preserve all specific measurements, colour values, and rules verbatim.

RAW REFERENCE MATERIAL:
{combined_markdown}

OUTPUT: A well-structured markdown specification document.""",

    "channel_best_practices": """You are a marketing channel specialist. Below is raw reference material about best practices for various marketing channels (social media, display, email, print, OOH) relevant to Barclays marketing.

Your task is to distil this into a clear, structured specification document that an AI channel compliance agent can use to review marketing proofs. Organise by channel type, then by platform/format. Focus on content guidelines, accessibility requirements, and engagement best practices.

Remove any redundancy or content not directly relevant to proof review. Preserve all specific recommendations and requirements.

RAW REFERENCE MATERIAL:
{combined_markdown}

OUTPUT: A well-structured markdown specification document.""",

    "channel_tech_specs": """You are a marketing production specialist. Below is raw reference material about technical specifications for various marketing channels (dimensions, file formats, file sizes, resolution requirements, platform constraints).

Your task is to distil this into a clear, structured specification document that an AI technical compliance agent can use to review marketing proofs. Organise by channel, then platform, then format. Use tables where appropriate for dimensions and specs.

Remove any redundancy or content not directly relevant to technical spec checking. Preserve all specific dimensions, file size limits, format requirements, and platform constraints verbatim.

RAW REFERENCE MATERIAL:
{combined_markdown}

OUTPUT: A well-structured markdown specification document.""",
}


class KnowledgeBaseService:
    """Orchestrates the document processing pipeline."""

    def __init__(self, llamaparse_service, gemini_service, reference_docs_service):
        self.llamaparse = llamaparse_service
        self.gemini = gemini_service
        self.reference_docs = reference_docs_service

    async def process_documents(
        self,
        kb_id: uuid.UUID,
        job_id: uuid.UUID,
        agent_key: str,
        user_id: Optional[uuid.UUID] = None,
        user_name: Optional[str] = None,
    ) -> None:
        """
        Run the full processing pipeline as a background task.

        1. Parse each source document with LlamaParse
        2. Combine parsed markdown
        3. Distil with Gemini into a spec
        4. Save as new SpecVersion
        5. Invalidate ReferenceDocsService cache
        """
        async with async_session_factory() as session:
            try:
                repo = KnowledgeBaseRepository(session)

                # Update job status to parsing
                await repo.update_processing_job(job_id, status="parsing_documents")
                await session.commit()

                # Get all source documents for this KB
                docs = await repo.get_source_documents(kb_id)
                if not docs:
                    await repo.update_processing_job(
                        job_id, status="failed",
                        error_message="No source documents found.",
                        completed_at=datetime.now(timezone.utc),
                    )
                    await session.commit()
                    return

                # Parse each document
                parsed_count = 0
                combined_parts = []
                source_doc_ids = []

                for doc in docs:
                    try:
                        # Load file data from storage
                        file_data = await storage_service.get_file(doc.file_storage_key)
                        if not file_data:
                            logger.error(f"[KB_SERVICE] File not found: {doc.file_storage_key}")
                            await repo.update_source_document_parse_status(
                                doc.id, "error", parse_error="File not found in storage."
                            )
                            await session.commit()
                            continue

                        # Parse with LlamaParse
                        await repo.update_source_document_parse_status(doc.id, "parsing")
                        await session.commit()

                        markdown = await self.llamaparse.parse_document(file_data, doc.filename)

                        # Update parse status
                        await repo.update_source_document_parse_status(
                            doc.id, "parsed", parsed_markdown=markdown
                        )
                        parsed_count += 1
                        await repo.update_processing_job(job_id, parsed_documents=parsed_count)
                        await session.commit()

                        if markdown.strip():
                            combined_parts.append(f"# {doc.filename}\n\n{markdown}")
                            source_doc_ids.append(str(doc.id))

                    except Exception as e:
                        logger.error(f"[KB_SERVICE] Error parsing {doc.filename}: {e}")
                        await repo.update_source_document_parse_status(
                            doc.id, "error", parse_error=str(e)
                        )
                        parsed_count += 1
                        await repo.update_processing_job(job_id, parsed_documents=parsed_count)
                        await session.commit()

                if not combined_parts:
                    await repo.update_processing_job(
                        job_id, status="failed",
                        error_message="No documents were successfully parsed.",
                        completed_at=datetime.now(timezone.utc),
                    )
                    await session.commit()
                    return

                # Distil with Gemini
                await repo.update_processing_job(job_id, status="distilling")
                await session.commit()

                combined_markdown = "\n\n---\n\n".join(combined_parts)
                prompt_template = DISTILLATION_PROMPTS.get(agent_key, DISTILLATION_PROMPTS["legal"])
                prompt = prompt_template.format(combined_markdown=combined_markdown)

                logger.info(f"[KB_SERVICE] Sending {len(combined_markdown)} chars to Gemini for distillation")
                response = await self.gemini.client.aio.models.generate_content(
                    model=self.gemini.model,
                    contents=prompt,
                )
                spec_content = response.text.strip()
                logger.info(f"[KB_SERVICE] Distillation complete: {len(spec_content)} chars")

                # Save as new spec version
                spec = await repo.create_spec_version(
                    knowledge_base_id=kb_id,
                    content=spec_content,
                    source_document_ids=source_doc_ids,
                    generated_by_id=user_id,
                    generated_by_name=user_name,
                    processing_job_id=job_id,
                )
                await session.commit()

                # Update job as completed
                await repo.update_processing_job(
                    job_id,
                    status="completed",
                    spec_version_id=spec.id,
                    completed_at=datetime.now(timezone.utc),
                )
                await session.commit()

                # Invalidate reference docs cache
                self.reference_docs.invalidate_cache(agent_key)
                logger.info(f"[KB_SERVICE] Pipeline complete for {agent_key}, spec version {spec.version_number}")

            except Exception as e:
                logger.error(f"[KB_SERVICE] Pipeline failed for job {job_id}: {e}")
                try:
                    await repo.update_processing_job(
                        job_id, status="failed",
                        error_message=str(e),
                        completed_at=datetime.now(timezone.utc),
                    )
                    await session.commit()
                except Exception:
                    logger.error("[KB_SERVICE] Failed to update job status after error")