python/src/agents/document_agent.py

"""
DocumentAgent - Conversational Document Management with PydanticAI

This agent enables users to create, update, and modify project documents through
natural conversation. Refactored for L2 modularity.
"""

import logging
import os
import uuid
from datetime import datetime
from typing import Any

from pydantic_ai import Agent, RunContext

from .base_agent import BaseAgent
from .document import logic
from .document.models import DocumentDependencies, DocumentOperation

logger = logging.getLogger(__name__)


class DocumentAgent(BaseAgent[DocumentDependencies, DocumentOperation]):
    """
    Conversational agent for document management.
    Facade pattern maintained for full backward compatibility.
    """

    def __init__(self, model: str | None = None, supabase_client: Any = None, **kwargs):
        if model is None:
            model = os.getenv("DOCUMENT_AGENT_MODEL")

        super().__init__(model=model, name="DocumentAgent", retries=3, enable_rate_limiting=True, **kwargs)
        self.supabase_client = supabase_client

    def _create_agent(self, **kwargs) -> Agent[DocumentDependencies, DocumentOperation]:
        """Create the PydanticAI agent with tools and prompts."""
        system_prompt = self.get_system_prompt()

        agent: Agent[DocumentDependencies, DocumentOperation] = Agent(
            model=self.model,
            deps_type=DocumentDependencies,
            system_prompt=system_prompt,
            **kwargs,
        )

        @agent.system_prompt
        async def add_project_context(ctx: RunContext[DocumentDependencies]) -> str:
            return f"""
**Current Project Context:**
- Project ID: {ctx.deps.project_id}
- User ID: {ctx.deps.user_id or "Unknown"}
- Current Document: {ctx.deps.current_document_id or "None"}
- Timestamp: {datetime.now().isoformat()}
"""

        @agent.tool
        async def list_documents(ctx: RunContext[DocumentDependencies]) -> str:
            """List all documents in the current project."""
            return await logic.list_documents_logic(self.supabase_client, ctx.deps.project_id)

        @agent.tool
        async def get_document(ctx: RunContext[DocumentDependencies], document_title: str) -> str:
            """Get the content of a specific document by title."""
            return await logic.get_document_logic(self.supabase_client, ctx.deps.project_id, document_title)

        @agent.tool
        async def create_document(
            ctx: RunContext[DocumentDependencies],
            title: str,
            document_type: str,
            content_description: str,
        ) -> str:
            """Create a new document with structured content."""
            return await logic.create_document_logic(
                ctx.deps.project_id,
                ctx.deps.user_id,
                title,
                document_type,
                content_description,
                ctx.deps.progress_callback,
            )

        @agent.tool
        async def update_document(
            ctx: RunContext[DocumentDependencies],
            document_title: str,
            section_to_update: str,
            new_content: str,
            update_description: str,
        ) -> str:
            """Update a specific section of an existing document."""
            return await logic.update_document_logic(
                ctx.deps.project_id, document_title, section_to_update, new_content, update_description
            )

        @agent.tool
        async def create_feature_plan(
            ctx: RunContext[DocumentDependencies],
            feature_name: str,
            feature_description: str,
            user_stories: str,
        ) -> str:
            """Create a React Flow feature plan."""
            return await logic.create_feature_plan_logic(
                ctx.deps.project_id, ctx.deps.user_id, feature_name, feature_description, user_stories
            )

        @agent.tool
        async def create_erd(
            ctx: RunContext[DocumentDependencies],
            system_name: str,
            entity_descriptions: str,
            relationships_description: str,
        ) -> str:
            """Create an Entity Relationship Diagram description and schema."""
            return await logic.create_erd_logic(
                ctx.deps.project_id, ctx.deps.user_id, system_name, entity_descriptions, relationships_description
            )

        @agent.tool
        async def request_approval(
            ctx: RunContext[DocumentDependencies],
            document_title: str,
            change_summary: str,
            change_type: str = "update",
        ) -> str:
            """Request approval for document changes."""
            return await logic.request_approval_logic(
                ctx.deps.project_id, ctx.deps.user_id, document_title, change_summary, change_type
            )

        return agent

    def get_system_prompt(self) -> str:
        """Fetch the system prompt from the prompt service."""
        default_prompt = """You are a Document Management Assistant.
You help users manage, create, and update project documentation.
You can list documents, create new ones, update specific sections, and generate diagrams like ERD or Feature Plans.
Always be professional and helpful."""
        try:
            from src.server.services.prompt_service import prompt_service

            prompt: str = prompt_service.get_prompt("document_agent_prompt", default_prompt)
            return prompt
        except (ImportError, Exception) as e:
            logger.warning(f"Could not load prompt from service (fallback to default): {e}")
            return default_prompt

    def _generate_block_id(self) -> str:
        """Generate a unique block ID (Facade)."""
        return str(uuid.uuid4())

    def _create_block(self, block_type: str, content: str, properties: dict[str, Any] | None = None) -> dict[str, Any]:
        """Create a document block (Facade)."""
        return {
            "id": self._generate_block_id(),
            "type": block_type,
            "content": content,
            "properties": properties or {"text": content},
        }

    async def run_conversation(
        self,
        user_message: str,
        project_id: str,
        user_id: str | None = None,
        current_document_id: str | None = None,
        progress_callback: Any | None = None,
    ) -> DocumentOperation:
        """Run the agent for conversational document management."""
        deps = DocumentDependencies(
            project_id=project_id,
            user_id=user_id,
            current_document_id=current_document_id,
            progress_callback=progress_callback,
        )

        try:
            result = await self.run(user_message, deps)
            return result
        except Exception as e:
            self.logger.error(f"Document operation failed: {str(e)}")
            return DocumentOperation(
                operation_type="error",
                document_id=None,
                document_type=None,
                title=None,
                success=False,
                message=f"Failed: {str(e)}",
                changes_made=[],
                content_preview=None,
            )