CrewAI Integration

from crewai import Crew, Agent, Task
from agent_os.integrations import crewai_kernel

# Define your agents as usual
researcher = Agent(
    role="Senior Research Analyst",
    goal="Discover groundbreaking insights about the topic",
    backstory="You are an expert researcher with deep analytical skills.",
    verbose=True,
    allow_delegation=True
)

writer = Agent(
    role="Content Writer",
    goal="Create engaging content based on research findings",
    backstory="You are a skilled writer who transforms research into compelling narratives.",
    verbose=True,
    allow_delegation=False
)

# Define tasks as usual
research_task = Task(
    description="Research the latest trends in {topic}",
    expected_output="A comprehensive research report with key findings",
    agent=researcher
)

writing_task = Task(
    description="Write an article based on the research findings",
    expected_output="A well-structured article ready for publication",
    agent=writer
)

# Create your crew as usual
crew = Crew(
    agents=[researcher, writer],
    tasks=[research_task, writing_task],
    verbose=True
)

# ✨ Wrap with Agent OS governance - ONE LINE
governed_crew = crewai_kernel.wrap(crew)

# Execute with full kernel protection
result = governed_crew.kickoff(inputs={"topic": "AI Safety"})
print(result)

from agent_os.integrations import crewai_kernel
from agent_os.protocols import IATPConfig, TrustLevel, CommunicationPolicy

# Configure Inter-Agent Trust Protocol
iatp_config = IATPConfig(
    # Trust levels between agent roles
    trust_matrix={
        ("Researcher", "Writer"): TrustLevel.HIGH,      # Researcher can share all findings
        ("Writer", "Researcher"): TrustLevel.MEDIUM,    # Writer can request clarifications
        ("Reviewer", "Writer"): TrustLevel.HIGH,        # Reviewer can send feedback
        ("Writer", "Reviewer"): TrustLevel.HIGH,        # Writer can submit for review
        ("Researcher", "Reviewer"): TrustLevel.LOW,     # Limited direct communication
    },
    
    # Default trust for unspecified pairs
    default_trust=TrustLevel.NONE,
    
    # Communication policies
    communication=CommunicationPolicy(
        # Message validation
        validate_messages=True,
        max_message_size_kb=100,
        
        # Content filtering
        filter_pii=True,
        filter_secrets=True,
        
        # Logging
        log_all_messages=True,
        log_destination="./logs/agent_comms.jsonl"
    )
)

# Configure what data can flow between agents
data_flow_rules = {
    # Researcher -> Writer: can share research, not raw data
    ("Researcher", "Writer"): {
        "allowed": ["summary", "findings", "citations"],
        "blocked": ["raw_data", "api_keys", "credentials"],
        "transform": {
            "pii": "redact",
            "numbers": "round_to_thousands"
        }
    },
    
    # Writer -> Reviewer: can share drafts
    ("Writer", "Reviewer"): {
        "allowed": ["draft", "outline", "metadata"],
        "blocked": ["internal_notes"]
    },
    
    # Reviewer -> Writer: can share feedback
    ("Reviewer", "Writer"): {
        "allowed": ["feedback", "approval_status", "revision_requests"],
        "blocked": ["other_submissions"]
    }
}

# Wrap crew with IATP
governed_crew = crewai_kernel.wrap(
    crew,
    trust_protocol="iatp",
    iatp_config=iatp_config,
    data_flow_rules=data_flow_rules,
    
    # Require mutual authentication before communication
    require_agent_auth=True,
    
    # Enable message signing for integrity
    sign_messages=True
)

"""
Governed Research Crew Example
==============================
A fully governed multi-agent research team with:
- Per-agent policies based on role
- Inter-agent trust protocol
- Task delegation governance
- Comprehensive audit logging
"""

from crewai import Crew, Agent, Task, Process
from crewai_tools import SerperDevTool, ScrapeWebsiteTool

from agent_os.integrations import crewai_kernel
from agent_os.policies import Policy, ResourcePolicy, ActionPolicy, DataPolicy
from agent_os.protocols import IATPConfig, TrustLevel, CommunicationPolicy
from agent_os.governance import DelegationPolicy, DelegationRule
from agent_os.audit import AuditConfig

# ============================================================
# 1. DEFINE AGENT POLICIES
# ============================================================

lead_researcher_policy = Policy(
    name="lead_researcher",
    resources=ResourcePolicy(
        file_access=["read"],
        network_access=["*.google.com", "*.arxiv.org", "*.github.com"],
        memory_limit_mb=1024
    ),
    actions=ActionPolicy(
        allowed_tools=["SerperDevTool", "ScrapeWebsiteTool", "analyze", "summarize"],
        max_tool_calls_per_task=100,
        max_tool_calls_per_minute=30
    ),
    data=DataPolicy(
        pii_handling="redact",
        output_sanitization=True
    )
)

data_analyst_policy = Policy(
    name="data_analyst",
    resources=ResourcePolicy(
        file_access=["read"],
        network_access=[],  # No external network access
        memory_limit_mb=2048  # More memory for analysis
    ),
    actions=ActionPolicy(
        allowed_tools=["analyze", "compute", "visualize"],
        blocked_tools=["search", "scrape"],  # Cannot fetch new data
        max_tool_calls_per_task=200
    )
)

writer_policy = Policy(
    name="writer",
    resources=ResourcePolicy(
        file_access=["read", "write"],
        allowed_paths=["./output/*", "./drafts/*"],
        network_access=[]
    ),
    actions=ActionPolicy(
        allowed_tools=["write_file", "format_markdown", "grammar_check"],
        blocked_tools=["search", "scrape", "execute_code"]
    )
)

reviewer_policy = Policy(
    name="reviewer",
    resources=ResourcePolicy(
        file_access=["read"],
        network_access=[]
    ),
    actions=ActionPolicy(
        allowed_tools=["review", "comment", "approve", "request_changes"],
        require_human_approval=["approve"]  # Final approval needs human
    )
)

# ============================================================
# 2. CONFIGURE INTER-AGENT TRUST
# ============================================================

iatp_config = IATPConfig(
    trust_matrix={
        ("Lead Researcher", "Data Analyst"): TrustLevel.HIGH,
        ("Data Analyst", "Lead Researcher"): TrustLevel.HIGH,
        ("Lead Researcher", "Technical Writer"): TrustLevel.HIGH,
        ("Data Analyst", "Technical Writer"): TrustLevel.MEDIUM,
        ("Technical Writer", "Quality Reviewer"): TrustLevel.HIGH,
        ("Quality Reviewer", "Technical Writer"): TrustLevel.HIGH,
    },
    default_trust=TrustLevel.LOW,
    communication=CommunicationPolicy(
        validate_messages=True,
        filter_pii=True,
        log_all_messages=True
    )
)

# ============================================================
# 3. CONFIGURE DELEGATION RULES
# ============================================================

delegation_policy = DelegationPolicy(
    allow_delegation=True,
    max_delegation_depth=2,
    rules=[
        DelegationRule(
            delegator="Lead Researcher",
            delegatee=["Data Analyst"],
            allowed_task_types=["analysis", "computation"]
        ),
        DelegationRule(
            delegator="Technical Writer",
            delegatee=[],  # Writer cannot delegate
            allowed_task_types=[]
        )
    ]
)

# ============================================================
# 4. CONFIGURE AUDIT LOGGING
# ============================================================

audit_config = AuditConfig(
    enabled=True,
    log_path="./audit_logs/research_crew.jsonl",
    log_level="INFO",
    include_timestamps=True,
    include_agent_context=True,
    include_tool_inputs=True,
    include_tool_outputs=True,
    redact_sensitive=True,  # Redact PII from logs
    rotation="daily",
    retention_days=90
)

# ============================================================
# 5. CREATE AGENTS
# ============================================================

search_tool = SerperDevTool()
scrape_tool = ScrapeWebsiteTool()

lead_researcher = Agent(
    role="Lead Researcher",
    goal="Conduct comprehensive research and coordinate the research team",
    backstory="""You are a senior research scientist with expertise in AI and ML.
    You lead research projects and coordinate with analysts and writers.""",
    tools=[search_tool, scrape_tool],
    allow_delegation=True,
    verbose=True
)

data_analyst = Agent(
    role="Data Analyst",
    goal="Analyze research data and extract meaningful insights",
    backstory="""You are an expert data analyst who transforms raw research 
    into actionable insights through statistical analysis.""",
    tools=[],
    allow_delegation=False,
    verbose=True
)

writer = Agent(
    role="Technical Writer",
    goal="Create clear, engaging technical content from research findings",
    backstory="""You are a skilled technical writer who makes complex topics
    accessible to a broad audience.""",
    tools=[],
    allow_delegation=False,
    verbose=True
)

reviewer = Agent(
    role="Quality Reviewer",
    goal="Ensure all content meets quality standards before publication",
    backstory="""You are a meticulous editor who ensures accuracy, clarity,
    and adherence to style guidelines.""",
    tools=[],
    allow_delegation=False,
    verbose=True
)

# ============================================================
# 6. DEFINE TASKS
# ============================================================

research_task = Task(
    description="""Research the topic: {topic}
    
    Requirements:
    1. Search for the latest developments and trends
    2. Identify key players and technologies
    3. Find relevant statistics and data points
    4. Compile a comprehensive research brief""",
    expected_output="A detailed research brief with sources and key findings",
    agent=lead_researcher
)

analysis_task = Task(
    description="""Analyze the research findings:
    
    1. Identify patterns and trends in the data
    2. Compare with historical context
    3. Provide statistical insights
    4. Highlight key takeaways""",
    expected_output="An analysis report with insights and visualizations",
    agent=data_analyst,
    context=[research_task]
)

writing_task = Task(
    description="""Write a comprehensive article:
    
    1. Create an engaging introduction
    2. Structure the content logically
    3. Include data visualizations
    4. Write a compelling conclusion""",
    expected_output="A polished article ready for review",
    agent=writer,
    context=[research_task, analysis_task]
)

review_task = Task(
    description="""Review the article for:
    
    1. Technical accuracy
    2. Clarity and readability
    3. Proper citations
    4. Grammar and style""",
    expected_output="A reviewed article with approval status",
    agent=reviewer,
    context=[writing_task]
)

# ============================================================
# 7. CREATE AND WRAP CREW
# ============================================================

crew = Crew(
    agents=[lead_researcher, data_analyst, writer, reviewer],
    tasks=[research_task, analysis_task, writing_task, review_task],
    process=Process.sequential,
    verbose=True
)

# ✨ Wrap with full governance
governed_crew = crewai_kernel.wrap(
    crew,
    
    # Per-agent policies
    agent_policies={
        "Lead Researcher": lead_researcher_policy,
        "Data Analyst": data_analyst_policy,
        "Technical Writer": writer_policy,
        "Quality Reviewer": reviewer_policy
    },
    
    # Inter-agent trust
    trust_protocol="iatp",
    iatp_config=iatp_config,
    
    # Delegation governance
    delegation_policy=delegation_policy,
    
    # Audit logging
    audit_config=audit_config,
    
    # Process isolation (optional - for high security)
    isolation="thread",  # or "process" for full isolation
    
    # Flight recorder for debugging
    flight_recorder=True
)

# ============================================================
# 8. EXECUTE
# ============================================================

if __name__ == "__main__":
    print("Starting governed research crew...")
    
    result = governed_crew.kickoff(
        inputs={"topic": "The Future of AI Agents in Enterprise"}
    )
    
    print("\n" + "="*60)
    print("FINAL OUTPUT")
    print("="*60)
    print(result)
    
    # Get governance metrics
    metrics = governed_crew.get_metrics()
    print("\n" + "="*60)
    print("GOVERNANCE METRICS")
    print("="*60)
    print(f"Total tool calls: {metrics['total_tool_calls']}")
    print(f"Blocked actions: {metrics['blocked_actions']}")
    print(f"Delegations: {metrics['total_delegations']}")
    print(f"Policy violations: {metrics['policy_violations']}")
    print(f"Audit entries: {metrics['audit_entries']}")