Add rate limiting error handling and update validation error types in FoundryKit

Author: rudrankriyam

.claude/settings.local.json

@@ -28,7 +28,8 @@
       "Bash(./find_unused_functions.sh)",
       "Bash(#!/bin/bash\n# Search for schema types\necho \"=== SEARCHING FOR SCHEMA TYPE USAGE ===\"\n\nsearch_type() {\n    local type_name=\"$1\"\n    echo \"\"\n    echo \"Searching for: $type_name\"\n    rg -c \"\\b$type_name\\b\" --type swift ./Sources ./Tests ./Examples 2>/dev/null | grep -v \":0$\" | sort\n}\n\n# Types from GenerationSchema.swift\nsearch_type \"SchemaType\"\nsearch_type \"RuntimeGenerationSchema\"\nsearch_type \"SchemaNode\"\nsearch_type \"Constraint\"\nsearch_type \"DynamicGenerationSchema\")",
       "Bash(swift package:*)",
-      "Bash(gh release edit:*)"
+      "Bash(gh release edit:*)",
+      "WebFetch(domain:github.com)"
     ],
     "deny": [],
     "includeCoAuthoredBy": false

Sources/FoundryKit/FoundationBackend.swift

@@ -217,6 +217,13 @@ internal final class FoundationBackend: FoundryBackend {
                     underlyingErrors: context.underlyingErrors
                 )
             )
+        case .rateLimited(let context):
+            return .rateLimited(
+                FoundryGenerationError.Context(
+                    debugDescription: context.debugDescription,
+                    underlyingErrors: context.underlyingErrors
+                )
+            )
         @unknown default:
             return .unknown(
                 FoundryGenerationError.Context(

Sources/FoundryKit/FoundryErrors.swift

@@ -48,6 +48,9 @@ public enum FoundryGenerationError: Error, LocalizedError {
     /// The requested feature is not supported.
     case unsupportedFeature(Context)
 
+    /// The request was rate limited.
+    case rateLimited(Context)
+    
     /// An unknown error occurred during generation.
     case unknown(Context)
 
@@ -73,6 +76,8 @@ public enum FoundryGenerationError: Error, LocalizedError {
             return "A network error occurred."
         case .unsupportedFeature:
             return "The requested feature is not supported."
+        case .rateLimited:
+            return "The request was rate limited. Please try again later."
         case .unknown:
             return "An unknown error occurred."
         }
@@ -100,6 +105,8 @@ public enum FoundryGenerationError: Error, LocalizedError {
             return "Check your internet connection and try again."
         case .unsupportedFeature:
             return "Use an alternative approach or check the documentation for supported features."
+        case .rateLimited:
+            return "Wait a moment before retrying or reduce the frequency of your requests."
         case .unknown:
             return "Try again or contact support if the problem persists."
         }
@@ -117,6 +124,7 @@ public enum FoundryGenerationError: Error, LocalizedError {
              .backendUnavailable(let context),
              .networkError(let context),
              .unsupportedFeature(let context),
+             .rateLimited(let context),
              .unknown(let context):
             return context.debugDescription
         }

Sources/FoundryKit/FoundryModelSession.swift

@@ -255,7 +255,7 @@ extension FoundryModelSession {
             let response = try await session.respond(to: prompt, schema: generationSchema)
 
             // Convert to our response type
-            let transcriptEntries = session.transcript.entries.suffix(2)
+            let transcriptEntries = Array(session.transcript).suffix(2)
             updateTranscript(with: transcriptEntries)
             return Response(content: response.content, transcriptEntries: transcriptEntries)
 
@@ -285,7 +285,7 @@ extension FoundryModelSession {
             let session = LanguageModelSession(model: .default, guardrails: guardrails, tools: tools)
             let response = try await session.respond(to: prompt, schema: generationSchema)
 
-            let transcriptEntries = session.transcript.entries.suffix(2)
+            let transcriptEntries = Array(session.transcript).suffix(2)
             updateTranscript(with: transcriptEntries)
             return Response(content: response.content, transcriptEntries: transcriptEntries)
 
@@ -461,9 +461,9 @@ extension FoundryModelSession {
 
     private func updateTranscript(with entries: ArraySlice<Transcript.Entry>) {
         // Update internal transcript with new entries
-        for entry in entries {
-            transcript.entries.append(entry)
-        }
+        // Note: In beta 2, Transcript.init(entries:) is private
+        // TODO: Find alternative approach for updating transcript
+        // For now, we'll rely on the session managing its own transcript
     }
 
     // Commented out for 0.0.1 release - focusing on simple text generation only

Sources/FoundryKit/Macros/FoundryGenerableMacros.swift

@@ -189,83 +189,83 @@ extension FoundryStructuredOutput {
             switch constraint {
             case .minimum(let minValue):
                 if let intValue = value as? Int, intValue < minValue {
-                    throw ValidationError.belowMinimum(name, minValue)
+                    throw MacroValidationError.belowMinimum(name, minValue)
                 }
             case .maximum(let maxValue):
                 if let intValue = value as? Int, intValue > maxValue {
-                    throw ValidationError.aboveMaximum(name, maxValue)
+                    throw MacroValidationError.aboveMaximum(name, maxValue)
                 }
             case .range(let range):
                 if let intValue = value as? Int, !range.contains(intValue) {
-                    throw ValidationError.outOfRange(name, range)
+                    throw MacroValidationError.outOfRange(name, range)
                 }
             case .minimumFloat(let minValue):
                 if let floatValue = value as? Float, floatValue < minValue {
-                    throw ValidationError.belowMinimumFloat(name, minValue)
+                    throw MacroValidationError.belowMinimumFloat(name, minValue)
                 }
             case .maximumFloat(let maxValue):
                 if let floatValue = value as? Float, floatValue > maxValue {
-                    throw ValidationError.aboveMaximumFloat(name, maxValue)
+                    throw MacroValidationError.aboveMaximumFloat(name, maxValue)
                 }
             case .rangeFloat(let range):
                 if let floatValue = value as? Float, !range.contains(floatValue) {
-                    throw ValidationError.outOfRangeFloat(name, range)
+                    throw MacroValidationError.outOfRangeFloat(name, range)
                 }
             case .minimumDouble(let minValue):
                 if let doubleValue = value as? Double, doubleValue < minValue {
-                    throw ValidationError.belowMinimumDouble(name, minValue)
+                    throw MacroValidationError.belowMinimumDouble(name, minValue)
                 }
             case .maximumDouble(let maxValue):
                 if let doubleValue = value as? Double, doubleValue > maxValue {
-                    throw ValidationError.aboveMaximumDouble(name, maxValue)
+                    throw MacroValidationError.aboveMaximumDouble(name, maxValue)
                 }
             case .rangeDouble(let range):
                 if let doubleValue = value as? Double, !range.contains(doubleValue) {
-                    throw ValidationError.outOfRangeDouble(name, range)
+                    throw MacroValidationError.outOfRangeDouble(name, range)
                 }
             case .minimumCount(let minCount):
                 if let arrayValue = value as? [Any], arrayValue.count < minCount {
-                    throw ValidationError.tooFewItems(name, minCount)
+                    throw MacroValidationError.tooFewItems(name, minCount)
                 }
             case .maximumCount(let maxCount):
                 if let arrayValue = value as? [Any], arrayValue.count > maxCount {
-                    throw ValidationError.tooManyItems(name, maxCount)
+                    throw MacroValidationError.tooManyItems(name, maxCount)
                 }
             case .count(let constraint):
                 if let arrayValue = value as? [Any] {
                     switch constraint {
                     case .exact(let exactCount):
                         if arrayValue.count != exactCount {
-                            throw ValidationError.wrongItemCount(name, exactCount, arrayValue.count)
+                            throw MacroValidationError.wrongItemCount(name, exactCount, arrayValue.count)
                         }
                     case .range(let range):
                         if !range.contains(arrayValue.count) {
-                            throw ValidationError.itemCountOutOfRange(name, range, arrayValue.count)
+                            throw MacroValidationError.itemCountOutOfRange(name, range, arrayValue.count)
                         }
                     }
                 }
             case .pattern(let regex):
                 if let stringValue = value as? String {
                     let predicate = NSPredicate(format: "SELF MATCHES %@", regex)
                     if !predicate.evaluate(with: stringValue) {
-                        throw ValidationError.patternMismatch(name, regex)
+                        throw MacroValidationError.patternMismatch(name, regex)
                     }
                 }
             case .anyOf(let allowed):
                 if let stringValue = value as? String, !allowed.contains(stringValue) {
-                    throw ValidationError.invalidEnumValue(name, stringValue, allowed)
+                    throw MacroValidationError.invalidEnumValue(name, stringValue, allowed)
                 }
             case .constant(let expected):
                 if let stringValue = value as? String, stringValue != expected {
-                    throw ValidationError.notEqualToConstant(name, expected, stringValue)
+                    throw MacroValidationError.notEqualToConstant(name, expected, stringValue)
                 }
             }
         }
     }
 }
 
 /// Validation errors
-internal enum ValidationError: LocalizedError {
+internal enum MacroValidationError: LocalizedError {
     case belowMinimum(String, Int)
     case aboveMaximum(String, Int)
     case outOfRange(String, ClosedRange<Int>)

Sources/FoundryKit/SimpleJSONValidator.swift

@@ -0,0 +1,153 @@
+// Commented out for 0.0.1 release - focusing on simple text generation only
+/*
+import Foundation
+import MLX
+import MLXLLM
+import MLXLMCommon
+import Tokenizers
+
+/// Simple JSON validation for generated content
+public struct SimpleJSONValidator {
+
+  /// Validate generated JSON against a schema with retry capability
+  public static func validateAndRetry<T: Codable>(
+    model: LanguageModel,
+    tokenizer: Tokenizer,
+    schema: [String: Any],
+    maxRetries: Int = 3,
+    prompt: String
+  ) async throws -> T {
+
+    for attempt in 1...maxRetries {
+      print("🔄 Validation attempt \(attempt)/\(maxRetries)")
+
+      // Generate JSON
+      let generatedText = try await generateText(
+        model: model,
+        tokenizer: tokenizer,
+        prompt: prompt
+      )
+
+      // Try to parse and validate
+      if let validObject: T = try? parseAndValidate(
+        jsonText: generatedText,
+        schema: schema
+      ) {
+        print("✅ Valid JSON generated on attempt \(attempt)")
+        return validObject
+      }
+
+      print("❌ Invalid JSON on attempt \(attempt), retrying...")
+    }
+
+    throw ValidationError.maxRetriesExceeded
+  }
+
+  /// Generate text using MLX model
+  private static func generateText(
+    model: LanguageModel,
+    tokenizer: Tokenizer,
+    prompt: String
+  ) async throws -> String {
+
+    let messages = [
+      ["role": "user", "content": prompt]
+    ]
+
+    let input = try await tokenizer.apply(chat: messages)
+    let parameters = GenerateParameters(maxTokens: 1000, temperature: 0.1)
+
+    var generatedText = ""
+    let generate = try await model.generate(
+      input: input,
+      parameters: parameters
+    ) { token in
+      generatedText += tokenizer.decode(tokens: [token])
+      return .more
+    }
+
+    return generatedText
+  }
+
+  /// Parse JSON and validate against schema
+  private static func parseAndValidate<T: Codable>(
+    jsonText: String,
+    schema: [String: Any]
+  ) throws -> T {
+
+    // Extract JSON from response (handle cases with extra text)
+    guard let cleanJSON = extractJSON(from: jsonText) else {
+      throw ValidationError.noJSONFound
+    }
+
+    // Basic JSON syntax validation
+    guard let data = cleanJSON.data(using: .utf8),
+      JSONSerialization.isValidJSONObject(try JSONSerialization.jsonObject(with: data))
+    else {
+      throw ValidationError.invalidJSONSyntax
+    }
+
+    // Decode to target type
+    let decoder = JSONDecoder()
+    do {
+      return try decoder.decode(T.self, from: data)
+    } catch {
+      throw ValidationError.schemaValidationFailed(error)
+    }
+  }
+
+  /// Extract JSON from potentially messy generated text
+  private static func extractJSON(from text: String) -> String? {
+    // Look for JSON blocks (between { and })
+    let patterns = [
+      "\\{[^{}]*(?:\\{[^{}]*\\}[^{}]*)*\\}",  // Simple nested objects
+      "```json\\s*([\\s\\S]*?)```",  // Markdown JSON blocks
+      "```\\s*([\\s\\S]*?)```",  // Generic code blocks
+    ]
+
+    for pattern in patterns {
+      if let range = text.range(of: pattern, options: .regularExpression) {
+        var jsonCandidate = String(text[range])
+
+        // Clean up markdown markers
+        jsonCandidate =
+          jsonCandidate
+          .replacingOccurrences(of: "```json", with: "")
+          .replacingOccurrences(of: "```", with: "")
+          .trimmingCharacters(in: .whitespacesAndNewlines)
+
+        return jsonCandidate
+      }
+    }
+
+    // Fallback: return the text as-is if it looks like JSON
+    let trimmed = text.trimmingCharacters(in: .whitespacesAndNewlines)
+    if trimmed.hasPrefix("{") && trimmed.hasSuffix("}") {
+      return trimmed
+    }
+
+    return nil
+  }
+}
+
+/// Validation errors
+public enum ValidationError: Error, LocalizedError {
+  case maxRetriesExceeded
+  case noJSONFound
+  case invalidJSONSyntax
+  case schemaValidationFailed(Error)
+
+  public var errorDescription: String? {
+    switch self {
+    case .maxRetriesExceeded:
+      return "Failed to generate valid JSON after maximum retries"
+    case .noJSONFound:
+      return "No JSON found in generated text"
+    case .invalidJSONSyntax:
+      return "Generated text is not valid JSON"
+    case .schemaValidationFailed(let error):
+      return "Schema validation failed: \(error.localizedDescription)"
+    }
+  }
+}
+*/