modules/ai_agent.py

import os
from langchain_bytez import BytezChatModel
from langchain.schema import HumanMessage, SystemMessage

def stream_compliance_report(summary_data, placeholder):
    """
    LangChain + Bytez implementation to stream an AI compliance report.
    """
    model = BytezChatModel(
        model_id="meta-llama/Llama-3.1-8B-Instruct",
        api_key=os.environ.get("BYTEZ_API_KEY"),
        max_new_tokens=1500,
        streaming=True
    )
    
    system_prompt = (
        "You are a Senior AML Compliance Analyst at a regulated "
        "financial institution. You write precise, formal reports "
        "reviewed by senior management and regulators. Always cite "
        "specific numbers. Never use casual language."
    )
    
    human_prompt = f"""
    Write a formal AML Compliance Monitoring Report with
    these exact sections:

    1. EXECUTIVE SUMMARY
       3-4 sentences covering analysis period,
       transactions reviewed, overall risk posture.

    2. KEY FINDINGS
       Bullet points of critical anomalies with
       exact numbers from the data.

    3. HIGH RISK TRANSACTIONS
       Describe high-risk patterns found: structuring
       attempts, large cash movements, international
       transfers. Reference specific counts.

    4. CUSTOMER RISK ASSESSMENT (KYC)
       Summarize KYC tier distribution. Flag customers
       with repeated suspicious behavior.

    5. REGULATORY IMPLICATIONS
       Reference BSA (Bank Secrecy Act), FinCEN SAR
       filing requirements, FATF Recommendation 16.
       State what filings or escalations are required.

    6. RECOMMENDATIONS
       Provide 5 specific, actionable recommendations
       for the compliance team.

    7. CONCLUSION
       Professional closing on AML posture and next steps.

    Data: {summary_data}

    Use formal regulatory language throughout.
    """
    
    messages = [
        SystemMessage(content=system_prompt),
        HumanMessage(content=human_prompt)
    ]
    
    full_report = ""
    try:
        for chunk in model.stream(messages):
            full_report += chunk.content
            placeholder.markdown(full_report + "▌")
        
        placeholder.markdown(full_report)
    except Exception as e:
        error_msg = f"Error generating report: {str(e)}"
        placeholder.error(error_msg)
        return error_msg
    finally:
        # Shutdown idle cluster
        if hasattr(model, 'shutdown_cluster'):
            model.shutdown_cluster()
            
    return full_report