Building an AI Contract Review Assistant

using System.Text.Json.Serialization;

public enum ClauseType
{
    Indemnification,
    LiabilityCap,
    Termination,
    IpOwnership,
    Confidentiality,
    PaymentTerms,
    ForceMajeure,
    GoverningLaw,
    DataProtection,
    NonCompete,
    Warranty,
    LimitationOfLiability
}

public enum RiskLevel
{
    Low,
    Medium,
    High,
    Critical
}

public record ExtractedClause(
    ClauseType ClauseType,
    string Text,
    int PageNumber,
    string SectionReference,
    RiskLevel? RiskLevel = null,
    string? RiskReason = null)
{
    public List<string> Recommendations { get; init; } = new();
}

public record Obligation(
    string Description,
    string Party,
    string? DueDate = null,
    bool IsRecurring = false,
    string SourceClause = "");

public record ContractState
{
    public required string DocumentId { get; init; }
    public required string Filename { get; init; }
    public required string ContractType { get; init; }
    public List<string> Parties { get; init; } = new();
    public string? EffectiveDate { get; init; }
    public string? ExpirationDate { get; init; }
    public List<DocumentSection> Sections { get; init; } = new();
    public List<ExtractedClause> ExtractedClauses { get; init; } = new();
    public List<Obligation> Obligations { get; init; } = new();
    public double? OverallRiskScore { get; init; }
    public string? RiskSummary { get; init; }
    public string? ExecutiveSummary { get; init; }
}

using Microsoft.SemanticKernel;

public class ContractReviewOrchestrator
{
    private readonly Kernel _kernel;
    private readonly IDocumentParserAgent _documentParser;
    private readonly IClauseExtractorAgent _clauseExtractor;
    private readonly IRiskAnalyzerAgent _riskAnalyzer;
    private readonly IObligationTrackerAgent _obligationTracker;
    private readonly IReportGeneratorAgent _reportGenerator;

    public ContractReviewOrchestrator(
        Kernel kernel,
        IDocumentParserAgent documentParser,
        IClauseExtractorAgent clauseExtractor,
        IRiskAnalyzerAgent riskAnalyzer,
        IObligationTrackerAgent obligationTracker,
        IReportGeneratorAgent reportGenerator)
    {
        _kernel = kernel;
        _documentParser = documentParser;
        _clauseExtractor = clauseExtractor;
        _riskAnalyzer = riskAnalyzer;
        _obligationTracker = obligationTracker;
        _reportGenerator = reportGenerator;
    }

    public async Task<ContractState> ReviewContractAsync(
        Stream documentStream,
        string filename,
        CancellationToken ct = default)
    {
        // Step 1: Parse document structure
        var state = await _documentParser.ParseAsync(
            documentStream, filename, ct);

        // Step 2: Extract and classify clauses
        state = await _clauseExtractor.ExtractAsync(state, ct);

        // Step 3: Parallel analysis
        var riskTask = _riskAnalyzer.AnalyzeAsync(state, ct);
        var obligationTask = _obligationTracker.TrackAsync(state, ct);

        await Task.WhenAll(riskTask, obligationTask);

        // Merge results
        state = state with
        {
            ExtractedClauses = riskTask.Result.ExtractedClauses,
            Obligations = obligationTask.Result.Obligations,
            OverallRiskScore = riskTask.Result.OverallRiskScore
        };

        // Step 4: Generate final report
        state = await _reportGenerator.GenerateAsync(state, ct);

        return state;
    }
}

using Azure;
using Azure.AI.DocumentIntelligence;

public class DocumentParserAgent : IDocumentParserAgent
{
    private readonly DocumentIntelligenceClient _client;
    private const int MaxChunkSize = 2000;
    private const int ChunkOverlap = 200;

    public DocumentParserAgent(string endpoint, string key)
    {
        _client = new DocumentIntelligenceClient(
            new Uri(endpoint),
            new AzureKeyCredential(key));
    }

    public async Task<ContractState> ParseAsync(
        Stream document,
        string filename,
        CancellationToken ct = default)
    {
        var operation = await _client.AnalyzeDocumentAsync(
            WaitUntil.Completed,
            "prebuilt-contract",
            document,
            cancellationToken: ct);

        var result = operation.Value;

        // Extract parties from document fields
        var parties = new List<string>();
        if (result.Documents?.FirstOrDefault()?.Fields
            .TryGetValue("Parties", out var partiesField) == true)
        {
            parties = partiesField.Value
                .AsArray()
                .Select(p => p.AsString())
                .ToList();
        }

        // Process sections with chunking
        var sections = new List<DocumentSection>();
        foreach (var paragraph in result.Paragraphs ?? Array.Empty<DocumentParagraph>())
        {
            var content = paragraph.Content;
            var pageNumber = paragraph.BoundingRegions?
                .FirstOrDefault()?.PageNumber ?? 1;

            if (content.Length > MaxChunkSize)
            {
                var chunks = ChunkWithOverlap(content);
                for (int i = 0; i < chunks.Count; i++)
                {
                    sections.Add(new DocumentSection(
                        Title: paragraph.Role ?? $"Section {sections.Count + 1}",
                        Content: chunks[i],
                        PageNumber: pageNumber,
                        ChunkIndex: i,
                        IsContinuation: i > 0
                    ));
                }
            }
            else
            {
                sections.Add(new DocumentSection(
                    Title: paragraph.Role ?? $"Section {sections.Count + 1}",
                    Content: content,
                    PageNumber: pageNumber,
                    ChunkIndex: 0,
                    IsContinuation: false
                ));
            }
        }

        return new ContractState
        {
            DocumentId = Guid.NewGuid().ToString(),
            Filename = filename,
            ContractType = DetectContractType(result.Content),
            Parties = parties,
            Sections = sections
        };
    }

    private List<string> ChunkWithOverlap(string text)
    {
        var chunks = new List<string>();
        var sentences = text.Split(new[] { ". ", ".\n" },
            StringSplitOptions.RemoveEmptyEntries);

        var currentChunk = new StringBuilder();
        foreach (var sentence in sentences)
        {
            if (currentChunk.Length + sentence.Length > MaxChunkSize)
            {
                chunks.Add(currentChunk.ToString().Trim());
                // Keep overlap from end of previous chunk
                var overlap = currentChunk.ToString()
                    .TakeLast(ChunkOverlap).ToString();
                currentChunk.Clear();
                currentChunk.Append(overlap);
            }
            currentChunk.Append(sentence).Append(". ");
        }

        if (currentChunk.Length > 0)
            chunks.Add(currentChunk.ToString().Trim());

        return chunks;
    }

    private string DetectContractType(string content)
    {
        var textLower = content.ToLower()[..Math.Min(5000, content.Length)];

        return textLower switch
        {
            var t when t.Contains("non-disclosure") => "NDA",
            var t when t.Contains("master service") => "MSA",
            var t when t.Contains("employment") && t.Contains("employee") => "Employment",
            var t when t.Contains("lease") && t.Contains("landlord") => "Lease",
            var t when t.Contains("software license") || t.Contains("saas") => "SaaS Agreement",
            _ => "General Contract"
        };
    }
}

using Microsoft.SemanticKernel;
using Microsoft.SemanticKernel.ChatCompletion;
using System.Text.Json;

public class ClauseExtractorAgent : IClauseExtractorAgent
{
    private readonly IChatCompletionService _chat;
    private const string ClassificationPrompt = """
        You are a legal contract analyst. Classify the following clause into ONE category.

        Categories:
        - indemnification: Party protects another from losses/claims
        - liability_cap: Limits on total damages or liability
        - termination: Conditions for ending contract, notice periods
        - ip_ownership: Ownership of intellectual property, work product
        - confidentiality: Protection of confidential information
        - payment_terms: Payment amounts, schedules, late fees
        - data_protection: GDPR, privacy, data handling requirements
        - force_majeure: Unforeseeable circumstances clauses
        - governing_law: Jurisdiction, venue for disputes
        - non_compete: Restrictions on competitive activities
        - warranty: Guarantees about quality, performance
        - limitation_of_liability: Exclusions of consequential damages
        - other: Does not fit any category

        Clause to classify:
        {0}

        Respond with JSON only:
        {{"clause_type": "", "confidence": <0.0-1.0>, "key_terms": ["term1", "term2"]}}
        """;

    public ClauseExtractorAgent(IChatCompletionService chat)
    {
        _chat = chat;
    }

    public async Task<ContractState> ExtractAsync(
        ContractState state,
        CancellationToken ct = default)
    {
        var extractedClauses = new List<ExtractedClause>();

        foreach (var section in state.Sections)
        {
            // Skip short sections (likely headers)
            if (section.Content.Length < 100)
                continue;

            var classification = await ClassifyClauseAsync(
                section.Content, ct);

            // Only keep high-confidence classifications
            if (classification.Confidence >= 0.7)
            {
                var clauseType = Enum.Parse<ClauseType>(
                    classification.ClauseType, ignoreCase: true);

                extractedClauses.Add(new ExtractedClause(
                    ClauseType: clauseType,
                    Text: section.Content,
                    PageNumber: section.PageNumber,
                    SectionReference: section.Title
                ));
            }
        }

        return state with { ExtractedClauses = extractedClauses };
    }

    private async Task<ClauseClassification> ClassifyClauseAsync(
        string text,
        CancellationToken ct)
    {
        var prompt = string.Format(ClassificationPrompt,
            text[..Math.Min(1500, text.Length)]);

        var settings = new PromptExecutionSettings
        {
            ExtensionData = new Dictionary<string, object>
            {
                ["temperature"] = 0,
                ["max_tokens"] = 100,
                ["response_format"] = new { type = "json_object" }
            }
        };

        var result = await _chat.GetChatMessageContentAsync(
            prompt, settings, cancellationToken: ct);

        return JsonSerializer.Deserialize<ClauseClassification>(
            result.Content)!;
    }
}

public record ClauseClassification(
    string ClauseType,
    double Confidence,
    List<string> KeyTerms
);

public class RiskAnalyzerAgent : IRiskAnalyzerAgent
{
    private readonly IChatCompletionService _chat;

    private const string RiskPrompt = """
        You are a legal risk analyst. Analyze this {0} clause for potential risks.

        Contract type: {1}
        Our position: {2}

        Clause text:
        {3}

        Evaluate against these criteria:
        1. Is this clause one-sided or balanced?
        2. Are there unusual or non-standard terms?
        3. What financial exposure does this create?
        4. Are there missing protections we should have?

        Risk levels:
        - low: Standard clause, balanced, minimal exposure
        - medium: Slightly unfavorable but common, manageable exposure
        - high: Significantly unfavorable, material exposure
        - critical: Extremely one-sided, major liability risk

        Respond with JSON:
        {{
            "risk_level": "",
            "risk_reason": "<2-3 sentence explanation>",
            "financial_exposure": "",
            "recommendations": ["", ""]
        }}
        """;

    public async Task<ContractState> AnalyzeAsync(
        ContractState state,
        CancellationToken ct = default)
    {
        var analyzedClauses = new List<ExtractedClause>();
        var riskScores = new List<int>();

        foreach (var clause in state.ExtractedClauses)
        {
            var result = await AnalyzeClauseRiskAsync(
                clause, state.ContractType, "customer", ct);

            var riskLevel = Enum.Parse<RiskLevel>(result.RiskLevel, true);

            analyzedClauses.Add(clause with
            {
                RiskLevel = riskLevel,
                RiskReason = result.RiskReason,
                Recommendations = result.Recommendations
            });

            riskScores.Add(RiskToScore(result.RiskLevel));
        }

        var overallScore = riskScores.Count > 0
            ? riskScores.Average() * 25
            : 0;

        return state with
        {
            ExtractedClauses = analyzedClauses,
            OverallRiskScore = overallScore,
            RiskSummary = GenerateRiskSummary(analyzedClauses)
        };
    }

    private int RiskToScore(string level) => level.ToLower() switch
    {
        "low" => 1,
        "medium" => 2,
        "high" => 3,
        "critical" => 4,
        _ => 2
    };

    private async Task<RiskAnalysisResult> AnalyzeClauseRiskAsync(
        ExtractedClause clause,
        string contractType,
        string ourPosition,
        CancellationToken ct)
    {
        var prompt = string.Format(RiskPrompt,
            clause.ClauseType,
            contractType,
            ourPosition,
            clause.Text[..Math.Min(2000, clause.Text.Length)]);

        var result = await _chat.GetChatMessageContentAsync(prompt,
            cancellationToken: ct);

        return JsonSerializer.Deserialize<RiskAnalysisResult>(result.Content)!;
    }
}

public record RiskAnalysisResult(
    string RiskLevel,
    string RiskReason,
    string FinancialExposure,
    List<string> Recommendations
);