diff --git a/FlagsmithClient/Classes/Flagsmith.swift b/FlagsmithClient/Classes/Flagsmith.swift
index 7a1a967..c7a7054 100644
--- a/FlagsmithClient/Classes/Flagsmith.swift
+++ b/FlagsmithClient/Classes/Flagsmith.swift
@@ -83,36 +83,52 @@ public final class Flagsmith: @unchecked Sendable {
     ///   - identity: ID of the user (optional)
     ///   - completion: Closure with Result which contains array of Flag objects in case of success or Error in case of failure
     public func getFeatureFlags(forIdentity identity: String? = nil,
+                                traits: [Trait]? = nil,
                                 completion: @Sendable @escaping (Result<[Flag], any Error>) -> Void)
     {
         if let identity = identity {
-            getIdentity(identity) { result in
-                switch result {
-                case let .success(thisIdentity):
-                    completion(.success(thisIdentity.flags))
-                case let .failure(error):
-                    if self.defaultFlags.isEmpty {
-                        completion(.failure(error))
-                    } else {
-                        completion(.success(self.defaultFlags))
+            if let traits = traits {
+                apiManager.request(.postTraits(identity: identity, traits: traits)) { (result: Result<Traits, Error>) in
+                    switch result {
+                    case let .success(result):
+                        completion(.success(result.flags))
+                    case let .failure(error):
+                        self.handleFlagsError(error, completion: completion)
+                    }
+                }
+            } else {
+                getIdentity(identity) { result in
+                    switch result {
+                    case let .success(thisIdentity):
+                        completion(.success(thisIdentity.flags))
+                    case let .failure(error):
+                        self.handleFlagsError(error, completion: completion)
                     }
                 }
             }
         } else {
-            apiManager.request(.getFlags) { (result: Result<[Flag], Error>) in
-                switch result {
-                case let .success(flags):
-                    completion(.success(flags))
-                case let .failure(error):
-                    if self.defaultFlags.isEmpty {
-                        completion(.failure(error))
-                    } else {
-                        completion(.success(self.defaultFlags))
+            if let _ = traits {
+                completion(.failure(FlagsmithError.invalidArgument("You must provide an identity to set traits")))
+            } else {
+                apiManager.request(.getFlags) { (result: Result<[Flag], Error>) in
+                    switch result {
+                    case let .success(flags):
+                        completion(.success(flags))
+                    case let .failure(error):
+                        self.handleFlagsError(error, completion: completion)
                     }
                 }
             }
         }
     }
+    
+    private func handleFlagsError(_ error: any Error, completion: @Sendable @escaping (Result<[Flag], any Error>) -> Void) {
+        if self.defaultFlags.isEmpty {
+            completion(.failure(error))
+        } else {
+            completion(.success(self.defaultFlags))
+        }
+    }
 
     /// Check feature exists and is enabled optionally for a specific identity
     ///
diff --git a/FlagsmithClient/Classes/FlagsmithError.swift b/FlagsmithClient/Classes/FlagsmithError.swift
index 2500e66..41403e1 100644
--- a/FlagsmithClient/Classes/FlagsmithError.swift
+++ b/FlagsmithClient/Classes/FlagsmithError.swift
@@ -21,6 +21,8 @@ public enum FlagsmithError: LocalizedError, Sendable {
     case decoding(DecodingError)
     /// Unknown or unhandled error was encountered.
     case unhandled(any Error)
+    /// Invalid argument error
+    case invalidArgument(String)
 
     public var errorDescription: String? {
         switch self {
@@ -36,6 +38,8 @@ public enum FlagsmithError: LocalizedError, Sendable {
             return "API Response could not be decoded: \(error.localizedDescription)"
         case let .unhandled(error):
             return "An unknown or unhandled error was encountered: \(error.localizedDescription)"
+        case let .invalidArgument(error):
+            return "Invalid argument error: \(error)"
         }
     }
 
diff --git a/FlagsmithClient/Classes/Traits.swift b/FlagsmithClient/Classes/Traits.swift
index da8d4ea..87ee8ec 100644
--- a/FlagsmithClient/Classes/Traits.swift
+++ b/FlagsmithClient/Classes/Traits.swift
@@ -13,4 +13,11 @@ import Foundation
 public struct Traits: Codable, Sendable {
     public let traits: [Trait]
     public let identifier: String?
+    public let flags: [Flag]
+    
+    init(traits: [Trait], identifier: String?, flags: [Flag] = []) {
+        self.traits = traits
+        self.identifier = identifier
+        self.flags = flags
+    }
 }
diff --git a/FlagsmithClient/Tests/RouterTests.swift b/FlagsmithClient/Tests/RouterTests.swift
index 45b2382..0d4846f 100644
--- a/FlagsmithClient/Tests/RouterTests.swift
+++ b/FlagsmithClient/Tests/RouterTests.swift
@@ -76,7 +76,8 @@ final class RouterTests: FlagsmithClientTestCase {
               "trait_value" : 42
             }
           ],
-          "identifier" : "A1B2C3D4"
+          "identifier" : "A1B2C3D4",
+          "flags": []
         }
         """.json(using: .utf8)
         let body = try request.httpBody.json()