Parameters:
 EnvironmentName:
    Type: String
    Default: LegalBlogSetup
    Description: Environment name for resource naming and tagging


Resources:
  InferenceProfile:
    Type: AWS::Bedrock::ApplicationInferenceProfile
    Properties:
      ModelSource:
        # Make sure to use the correct model ARN for your region
        CopyFrom: !Sub "arn:aws:bedrock:${AWS::Region}:${AWS::AccountId}:inference-profile/us.anthropic.claude-3-5-haiku-20241022-v1:0"
      InferenceProfileName: "Haiku35-AgentProfile"
      Description: "Inference profile for Haiku 3.5 v1"
  
  # Bedrock Agent Execution Role
  BedrockAgentExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: bedrock.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: BedrockAgentCustomPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Sid: AmazonBedrockAgentInferenceProfilesCrossRegionPolicyProd
                Effect: Allow
                Action:
                  - bedrock:InvokeModel
                  - bedrock:InvokeModelWithResponseStream
                  - bedrock:GetInferenceProfile
                  - bedrock:GetFoundationModel
                Resource:
                  - !GetAtt InferenceProfile.InferenceProfileArn
                  - "arn:aws:bedrock:*::foundation-model/anthropic.claude-3-5-haiku-20241022-v1:0"

  #Classification Agent and Alias
  ClassificationAgent:
    Type: AWS::Bedrock::Agent
    Properties:
      AgentName: !Sub "${EnvironmentName}-Classification-Agent"
      Description: Agent for Document Classification & PII-Detection
      AutoPrepare: true
      Instruction: |
        Step 1) Classify this document into one of these categories:
                  - Contract
                  - Email
                  - Legal
        
        Step 2) Identify any PII (Personal Identifiable Information) in the following text:
                  - Social Security Numbers
                  - Email addresses
                  - Phone numbers
                  - Physical addresses

        Document: {text}
        
        Return JSON format:
        {{
            "category": "category_name",
            "confidence_score": 0-100,
            "reasoning": "brief explanation"
        }},
        {{
            "found_pii": {{
                "ssn": [],
                "emails": [],
                "phones": [],
                "addresses": []
            }},
            "risk_level": "HIGH/MEDIUM/LOW",
            "total_pii_count": 0
        }}
      AgentResourceRoleArn: !GetAtt BedrockAgentExecutionRole.Arn
      IdleSessionTTLInSeconds: 1800
      FoundationModel: !GetAtt InferenceProfile.InferenceProfileArn

  ClassificationAgentAlias:
    Type: AWS::Bedrock::AgentAlias
    Properties:
      AgentAliasName: !Sub "${EnvironmentName}-ClassificationAgent"
      AgentId: !Ref ClassificationAgent

  #Contract Agent and Alias
  ContractAgent:
    Type: AWS::Bedrock::Agent
    Properties:
      AgentName: !Sub "${EnvironmentName}-Contract-Analysis-Agent"
      Description: Agent for analyzing contracts
      AutoPrepare: true
      Instruction: |
        Analyze this contract for:
        - Parties involved
        - Key dates
        - Financial terms
        - Obligations
        - Termination clauses
        
        Contract: {text}
        
        Return JSON format:
        {{
            "parties": [],
            "key_dates": {{
                "start_date": "",
                "end_date": "",
                "other_important_dates": []
            }},
            "financial_terms": {{
                "amounts": [],
                "payment_schedule": "",
                "currency": ""
            }},
            "obligations": {{
                "party_1": [],
                "party_2": []
            }},
            "termination": {{
                "conditions": [],
                "notice_period": "",
                "penalties": []
            }}
        }}
      AgentResourceRoleArn: !GetAtt BedrockAgentExecutionRole.Arn
      IdleSessionTTLInSeconds: 1800
      FoundationModel: !GetAtt InferenceProfile.InferenceProfileArn

  ContractAgentAlias:
    Type: AWS::Bedrock::AgentAlias
    Properties:
      AgentAliasName: !Sub "${EnvironmentName}-ContractAnalysisAgent"
      AgentId: !Ref ContractAgent

  #Email Agent and Alias
  EmailAgent:
    Type: AWS::Bedrock::Agent
    Properties:
      AgentName: !Sub "${EnvironmentName}-Email-Analysis-Agent"
      Description: Agent for analyzing Email communications
      AutoPrepare: true
      Instruction: |
        Analyze this email for:
        - Sender/Recipients
        - Key topics
        - Action items
        - Sentiment
        - Potential legal implications
        
          Email: {text}
        
        Return JSON format:
        {{
            "participants": {{
                "sender": "",
                "to": [],
                "cc": [],
                "bcc": []
            }},
            "topics": [],
            "action_items": [],
            "sentiment": {{
                "type": "POSITIVE/NEUTRAL/NEGATIVE",
                "score": 0-100
            }},
            "legal_implications": []
        }}
      AgentResourceRoleArn: !GetAtt BedrockAgentExecutionRole.Arn
      IdleSessionTTLInSeconds: 1800
      FoundationModel: !GetAtt InferenceProfile.InferenceProfileArn

  EmailAgentAlias:
    Type: AWS::Bedrock::AgentAlias
    Properties:
      AgentAliasName: !Sub "${EnvironmentName}-EmailAnalysisAgent"
      AgentId: !Ref EmailAgent

  #Legal Agent and Alias
  LegalAgent:
    Type: AWS::Bedrock::Agent
    Properties:
      AgentName: !Sub "${EnvironmentName}-Legal-Analysis-Agent"
      Description: Agent for analyzing Legal Documents
      AutoPrepare: true
      Instruction: |
        Analyze for:
        - Legal arguments
        - Citations
        - Precedents
        - Risk factors
        - Privilege indicators
        
        Text: {text}
        
        Return JSON format:
        {{
            "legal_arguments": [],
            "citations": [],
            "precedents": [],
            "risk_factors": {{
                "high": [],
                "medium": [],
                "low": []
            }},
            "privilege": {{
                "is_privileged": true/false,
                "type": "",
                "reasoning": ""
            }}
        }}
      AgentResourceRoleArn: !GetAtt BedrockAgentExecutionRole.Arn
      IdleSessionTTLInSeconds: 1800
      FoundationModel: !GetAtt InferenceProfile.InferenceProfileArn

  LegalAgentAlias:
    Type: AWS::Bedrock::AgentAlias
    Properties:
      AgentAliasName: !Sub "${EnvironmentName}-LegalAnalysisAgent"
      AgentId: !Ref LegalAgent

  # IAM Policy for the collaborative Bedrock Agent
  CollabAgentPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Description: "Policy for Collaborative Bedrock Agent to invoke foundation models and collaborate with other agents"
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Sid: "AmazonBedrockAgentBedrockFoundationModelPolicyProd"
            Effect: "Allow"
            Action:
              - "bedrock:InvokeModel"
              - "bedrock:InvokeModelWithResponseStream"
              - "bedrock:GetInferenceProfile"
              - "bedrock:GetFoundationModel"
            Resource:
              - !GetAtt InferenceProfile.InferenceProfileArn
              - "arn:aws:bedrock:*::foundation-model/anthropic.claude-3-5-haiku-20241022-v1:0"
          
          - Sid: "AmazonBedrockAgentsMultiAgentsPoliciesProd"
            Effect: "Allow"
            Action:
              - "bedrock:GetAgentAlias"
              - "bedrock:InvokeAgent"
            Resource:
              - !GetAtt LegalAgentAlias.AgentAliasArn
              - !GetAtt EmailAgentAlias.AgentAliasArn
              - !GetAtt ContractAgentAlias.AgentAliasArn
              - !GetAtt ClassificationAgentAlias.AgentAliasArn

  # IAM Role for the collaborator agent
  BedrockCollabAgentExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: bedrock.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - !Ref CollabAgentPolicy
      Description: "Role for Collaborative Bedrock Agent to invoke foundation models and collaborate with other agents"

  CollabBedrockAgent:
    Type: AWS::Bedrock::Agent
    Properties:
      AgentName: !Sub "${EnvironmentName}-eDiscoveryCollaboratorAgent"
      FoundationModel: !GetAtt InferenceProfile.InferenceProfileArn
      Instruction: |
          You are a supervisor agent that coordinates document analysis across specialized agents. Your primary responsibility is to orchestrate the document analysis workflow by:
            1. First using the Document Classification Agent to identify document type and detect sensitive PII
            2. Then routing to the appropriate specialized agent based on document classification
            3. Finally synthesizing all insights into a comprehensive, actionable report

          Available Specialized Agents
            You will coordinate with these Amazon Bedrock agents:
              1. Document Classification Agent - Identifies document type and flags sensitive PII data
              2. Contracts Agent - Reviews and extracts key elements from contracts, agreements, and terms
              3. Email Analysis Agent - Processes email communications, threads, and correspondence
              4. Legal Document Agent - Analyzes legal documentation, compliance issues, and implications

          Workflow Process
            1. Initial Classification:
              - Submit document to Document Classification Agent
              - Receive document type and PII assessment
              - Determine appropriate specialized agent

            2. Specialized Analysis:
              - Route document to the appropriate specialized agent
              - Request comprehensive analysis based on document type
              - Receive detailed findings in JSON format

            3. Synthesis and Reporting:
              - Integrate all findings from specialized agents
              - Generate comprehensive analysis report
              - Highlight critical insights and recommended actions

          Comprehensive Report Generation
            After receiving the specialized agent's complete analysis, synthesize all findings into a comprehensive report:

            1. Agent Consultation Record:
                Agent Consulted: [Contract/Email/Legal Document] Agent
                Agent Response: [INCLUDE FULL JSON RESPONSE]

            2. Final Analysis Report:
                {
                  "executive_summary": {
                    "key_points": [],
                    "risk_level": "HIGH/MEDIUM/LOW",
                    "immediate_actions": []
                  },
                  "key_findings": {
                    "pii_related": [],
                    "legal_related": [],
                    "business_related": []
                  },
                  "risk_assessment": {
                    "overall_risk": "HIGH/MEDIUM/LOW",
                    "risk_factors": [],
                    "mitigation_suggestions": []
                  },
                  "recommended_actions": {
                    "immediate": [],
                    "short_term": [],
                    "long_term": []
                  },
                  "detailed_analysis": {
                    "by_category": {},
                    "cross_references": [],
                    "special_considerations": []
                  }
                }

          Response Requirements
            - Always maintain document confidentiality and security
            - Prioritize findings based on risk level and business impact
            - Ensure all PII concerns are prominently highlighted
            - Provide actionable recommendations with clear justification
            - Format all outputs according to the specified JSON structure
      AgentResourceRoleArn: !GetAtt BedrockCollabAgentExecutionRole.Arn
      IdleSessionTTLInSeconds: 1800
      AutoPrepare: true
      AgentCollaboration: "SUPERVISOR"
      AgentCollaborators:
        - AgentDescriptor:
            AliasArn: !GetAtt LegalAgentAlias.AgentAliasArn
          CollaborationInstruction: |
            Analyze for:
              - Legal arguments
              - Citations
              - Precedents
              - Risk factors
              - Privilege indicators
              
              Return JSON format:
              {{
                  "legal_arguments": [],
                  "citations": [],
                  "precedents": [],
                  "risk_factors": {{
                      "high": [],
                      "medium": [],
                      "low": []
                  }},
                  "privilege": {{
                      "is_privileged": true/false,
                      "type": "",
                      "reasoning": ""
                  }}
              }}
          CollaboratorName: "LegalAgent"
          RelayConversationHistory: "TO_COLLABORATOR"
        
        - AgentDescriptor:
            AliasArn: !GetAtt EmailAgentAlias.AgentAliasArn
          CollaborationInstruction: |
              Analyze this email for:
                - Sender/Recipients
                - Key topics
                - Action items
                - Sentiment
                - Potential legal implications
                
                Return JSON format:
                {{
                    "participants": {{
                        "sender": "",
                        "to": [],
                        "cc": [],
                        "bcc": []
                    }},
                    "topics": [],
                    "action_items": [],
                    "sentiment": {{
                        "type": "POSITIVE/NEUTRAL/NEGATIVE",
                        "score": 0-100
                    }},
                    "legal_implications": []
                }}
          CollaboratorName: "EmailAgent"
          RelayConversationHistory: "TO_COLLABORATOR"

        - AgentDescriptor:
            AliasArn: !GetAtt ContractAgentAlias.AgentAliasArn
          CollaborationInstruction: |
              Analyze this contract for:
              - Parties involved
              - Key dates
              - Financial terms
              - Obligations
              - Termination clauses
              
              Return JSON format:
              {{
                  "parties": [],
                  "key_dates": {{
                      "start_date": "",
                      "end_date": "",
                      "other_important_dates": []
                  }},
                  "financial_terms": {{
                      "amounts": [],
                      "payment_schedule": "",
                      "currency": ""
                  }},
                  "obligations": {{
                      "party_1": [],
                      "party_2": []
                  }},
                  "termination": {{
                      "conditions": [],
                      "notice_period": "",
                      "penalties": []
                  }}
              }}
          CollaboratorName: "ContractAgent"
          RelayConversationHistory: "TO_COLLABORATOR"

        - AgentDescriptor:
            AliasArn: !GetAtt ClassificationAgentAlias.AgentAliasArn
          CollaborationInstruction: |
            Step 1) Classify this document into one of these categories:
                      - Contract
                      - Email
                      - Legal
            
            Step 2) Identify any PII (Personal Identifiable Information) in the following text:
                      - Social Security Numbers
                      - Email addresses
                      - Phone numbers
                      - Physical addresses
            
            Return JSON format:
            {{
                "category": "category_name",
                "confidence_score": 0-100,
                "reasoning": "brief explanation"
            }},
            {{
                "found_pii": {{
                    "ssn": [],
                    "emails": [],
                    "phones": [],
                    "addresses": []
                }},
                "risk_level": "HIGH/MEDIUM/LOW",
                "total_pii_count": 0
            }}
          CollaboratorName: "ClassificationAgent"
          RelayConversationHistory: "TO_COLLABORATOR"

  CollabBedrockAgentAlias:
    Type: AWS::Bedrock::AgentAlias
    Properties:
      AgentAliasName: !Sub "${EnvironmentName}-ACollabBedrockAgent"
      AgentId: !Ref CollabBedrockAgent

Outputs:
  CollabBedrockAgentId:
    Description: Collaborator Agent ID
    Value: !GetAtt CollabBedrockAgent.AgentId
    Export:
      Name: !Sub "${EnvironmentName}-CollabBedrockAgentId"

  CollabBedrockAgentAliasId:
    Description: Collaborator Agent Alias ID 
    Value: !GetAtt CollabBedrockAgentAlias.AgentAliasId
    Export:
      Name: !Sub "${EnvironmentName}-CollabBedrockAgentAliasId"
