Data Extraction Pipeline

import { z } from "zod";

const ContactInfo = z.object({
  name: z.string().describe("Full name of the person"),
  email: z.string().email().optional().describe("Email address if present"),
  phone: z.string().optional().describe("Phone number if present"),
  company: z.string().optional().describe("Company or organization name"),
  role: z.string().optional().describe("Job title or role"),
});

const ExtractedTicket = z.object({
  subject: z.string().describe("Brief summary of the issue"),
  priority: z.enum(["low", "medium", "high", "critical"]),
  category: z.enum(["billing", "technical", "account", "feature_request", "other"]),
  contact: ContactInfo,
  sentiment: z.enum(["positive", "neutral", "negative", "frustrated"]),
  actionItems: z.array(z.string()).describe("Concrete next steps to resolve"),
});

import type { OutputGuardrail } from "stratus-sdk/core";

const extractionQualityGuard: OutputGuardrail = {
  name: "extraction_quality",
  execute: (output) => {
    try {
      const data = JSON.parse(output);

      // Reject if no action items were generated
      if (!data.actionItems || data.actionItems.length === 0) {
        return {
          tripwireTriggered: true,
          outputInfo: { reason: "No action items extracted" },
        };
      }

      // Reject if contact has no name
      if (!data.contact?.name) {
        return {
          tripwireTriggered: true,
          outputInfo: { reason: "Contact name is required" },
        };
      }

      return { tripwireTriggered: false };
    } catch {
      return {
        tripwireTriggered: true,
        outputInfo: { reason: "Invalid JSON output" },
      };
    }
  },
};

const piiRedactionGuard: OutputGuardrail = {
  name: "pii_check",
  execute: (output) => {
    // Check for SSNs, credit card numbers, etc. that shouldn't be in extraction output
    const hasSensitivePII = /\b\d{3}-\d{2}-\d{4}\b/.test(output) ||
                            /\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b/.test(output);
    return {
      tripwireTriggered: hasSensitivePII,
      outputInfo: { reason: "Sensitive PII detected in extraction output" },
    };
  },
};

import { prompt } from "stratus-sdk/core";

async function extractFromDocuments(documents: string[]) {
  const results: z.infer<typeof ExtractedTicket>[] = [];

  for (const doc of documents) {
    try {
      const result = await prompt(doc, {
        model,
        instructions: `Extract structured data from the following support ticket.`,
        outputType: ExtractedTicket,
        outputGuardrails: [extractionQualityGuard],
      });
      results.push(result.finalOutput);
    } catch (error) {
      if (error instanceof OutputGuardrailTripwireTriggered) {
        console.warn(`Skipped document: ${error.outputInfo}`);
      } else {
        throw error;
      }
    }
  }

  return results;
}

import { withTrace } from "stratus-sdk/core";

const { result, trace } = await withTrace("ticket_extraction", () =>
  run(enrichedExtractor, ticketText)
);

console.log(`Extraction took ${trace.duration}ms`);
console.log(`Model calls: ${trace.spans.filter(s => s.type === "model_call").length}`);
console.log(`Tool calls: ${trace.spans.filter(s => s.type === "tool_execution").length}`);
console.log(`Priority: ${result.finalOutput.priority}`);