diff --git a/README.md b/README.md
index 1d27c31..cf73a3b 100644
--- a/README.md
+++ b/README.md
@@ -34,6 +34,23 @@ yarn:
 yarn add @ftrack/api
 ```
 
+## TypeScript: Generated Schema Types
+
+You can generate schema types for your own workspace with [@ftrack/ts-schema-generator](https://github.com/ftrackhq/ftrack-ts-schema-generator).
+
+Once generated, you can integrate them with @ftrack/api by passing them as a type variable:
+
+```ts
+import SchemaTypes from "./__generated__/schema.ts"
+import { Session } from @ftrack/api
+
+const session = new Session<SchemaTypes>(...);
+
+// user will now be of type User
+// and provide all available properties for its entity type.
+const user = await session.query<"User">("...");
+```
+
 ## Tutorial
 
 The API uses sessions to manage communication with an ftrack server. Create a session that connects to your ftrack server (changing the passed values as appropriate):
diff --git a/source/operation.ts b/source/operation.ts
index af40637..ca295f8 100644
--- a/source/operation.ts
+++ b/source/operation.ts
@@ -4,9 +4,9 @@
  * @namespace operation
  */
 
-export interface CreateOperation {
+export interface CreateOperation<TEntityType> {
   action: "create";
-  entity_type: string;
+  entity_type: TEntityType;
   entity_data: any;
 }
 
@@ -15,33 +15,33 @@ export interface QueryOperation {
   expression: string;
 }
 
-export interface SearchOperationOptions {
+export interface SearchOperationOptions<TEntityType> {
   expression?: string;
-  entityType?: string;
+  entityType?: TEntityType;
   terms?: string[];
   contextId?: string;
   objectTypeIds?: string[];
 }
 
-export interface SearchOperation {
+export interface SearchOperation<TEntityType> {
   action: "search";
   expression?: string;
-  entity_type?: string;
+  entity_type?: TEntityType;
   terms?: string[];
   context_id?: string;
   object_type_ids?: string[];
 }
 
-export interface UpdateOperation {
+export interface UpdateOperation<TEntityType> {
   action: "update";
-  entity_type: string;
+  entity_type: TEntityType;
   entity_key: string[] | string;
   entity_data: any;
 }
 
-export interface DeleteOperation {
+export interface DeleteOperation<TEntityType> {
   action: "delete";
-  entity_type: string;
+  entity_type: TEntityType;
   entity_key: string[] | string;
 }
 
@@ -61,12 +61,12 @@ export interface GetUploadMetadataOperation {
   component_id: string;
 }
 
-export type Operation =
-  | CreateOperation
+export type Operation<TEntityType> =
+  | CreateOperation<TEntityType>
   | QueryOperation
-  | SearchOperation
-  | UpdateOperation
-  | DeleteOperation
+  | SearchOperation<TEntityType>
+  | UpdateOperation<TEntityType>
+  | DeleteOperation<TEntityType>
   | QueryServerInformationOperation
   | QuerySchemasOperation
   | GetUploadMetadataOperation
@@ -81,7 +81,10 @@ export type Operation =
  * @param  {Object} data Entity data to use for creation
  * @return {Object}      API operation
  */
-export function create(type: string, data: any): CreateOperation {
+export function create<TEntityType>(
+  type: TEntityType,
+  data: any,
+): CreateOperation<TEntityType> {
   return {
     action: "create",
     entity_type: type,
@@ -109,13 +112,13 @@ export function query(expression: string): QueryOperation {
  * @param  {string} expression API query expression
  * @return {Object}            API operation
  */
-export function search({
+export function search<TEntityType>({
   expression,
   entityType,
   terms,
   contextId,
   objectTypeIds,
-}: SearchOperationOptions): SearchOperation {
+}: SearchOperationOptions<TEntityType>): SearchOperation<TEntityType> {
   return {
     action: "search",
     expression,
@@ -136,11 +139,11 @@ export function search({
  * @param  {Object} data values to update
  * @return {Object}      API operation
  */
-export function update(
-  type: string,
+export function update<TEntityType>(
+  type: TEntityType,
   keys: string[] | string,
   data: any,
-): UpdateOperation {
+): UpdateOperation<TEntityType> {
   return {
     action: "update",
     entity_type: type,
@@ -158,10 +161,10 @@ export function update(
  * @param  {Array} keys Identifying keys, typically [<entity id>]
  * @return {Object}      API operation
  */
-function deleteOperation(
-  type: string,
+function deleteOperation<TEntityType>(
+  type: TEntityType,
   keys: string[] | string,
-): DeleteOperation {
+): DeleteOperation<TEntityType> {
   return {
     action: "delete",
     entity_type: type,
diff --git a/source/project_schema.ts b/source/project_schema.ts
index 7105c95..b15a2da 100644
--- a/source/project_schema.ts
+++ b/source/project_schema.ts
@@ -2,7 +2,7 @@
 
 import * as operation from "./operation.js";
 import { Session } from "./session.js";
-import type { Data, QueryResponse } from "./types.js";
+import type { QueryResponse } from "./types.js";
 /**
  * Project schema namespace
  * @namespace project_schema
@@ -18,8 +18,8 @@ import type { Data, QueryResponse } from "./types.js";
  *
  * @memberof project_schema
  */
-export function getStatuses(
-  session: Session,
+export function getStatuses<TEntityTypeMap extends Record<string, any>>(
+  session: Session<TEntityTypeMap>,
   projectSchemaId: string,
   entityType: string,
   typeId: string | null = null,
@@ -71,7 +71,8 @@ export function getStatuses(
     ),
   );
 
-  response = session.call<QueryResponse<Data>>(operations);
+  response =
+    session.call<QueryResponse<TEntityTypeMap["ProjectSchema"]>>(operations);
   response = response.then((results) => {
     // Since the operations where performed in one batched call,
     // the result will be merged into a single entity.
diff --git a/source/session.ts b/source/session.ts
index 4fe3e37..ca06f57 100644
--- a/source/session.ts
+++ b/source/session.ts
@@ -18,8 +18,8 @@ import type {
   CreateComponentOptions,
   CreateResponse,
   Data,
+  DefaultEntityTypeMap,
   DeleteResponse,
-  Entity,
   GetUploadMetadataResponse,
   IsTuple,
   MutationOptions,
@@ -48,7 +48,9 @@ const ENCODE_DATETIME_FORMAT = "YYYY-MM-DDTHH:mm:ss";
  * @class  Session
  *
  */
-export class Session {
+export class Session<
+  TEntityTypeMap extends Record<string, any> = DefaultEntityTypeMap,
+> {
   apiUser: string;
   apiKey: string;
   serverUrl: string;
@@ -56,17 +58,19 @@ export class Session {
   eventHub: EventHub;
   clientToken: string | null;
   initialized: boolean;
-  initializing: Promise<Session>;
+  initializing: Promise<Session<TEntityTypeMap>>;
   additionalHeaders: Data;
-  schemas?: Schema[];
+  schemas?: Schema<TEntityTypeMap>[];
   serverInformation?: QueryServerInformationResponse;
   serverVersion?: string;
   private ensureSerializableResponse: boolean;
   private decodeDatesAsIso: boolean;
-  private schemasPromise?: Promise<Schema[]>;
+  private schemasPromise?: Promise<Schema<TEntityTypeMap>[]>;
   private serverInformationPromise?: Promise<ServerInformation>;
   private serverInformationValues?: string[];
-  private schemaMapping?: Record<string, Schema>;
+  private schemaMapping?: {
+    [TEntityType in keyof TEntityTypeMap]: Schema<TEntityTypeMap, TEntityType>;
+  };
 
   /**
    * Construct Session instance with API credentials.
@@ -187,7 +191,10 @@ export class Session {
     }
 
     // TODO: remove these operations from session initialization in next major
-    const operations: operation.Operation[] = [
+    const operations: [
+      operation.QueryServerInformationOperation,
+      operation.QuerySchemasOperation,
+    ] = [
       {
         action: "query_server_information",
         values: serverInformationValues,
@@ -216,9 +223,9 @@ export class Session {
     this.initialized = false;
 
     const initializingPromise =
-      this.call<[QueryServerInformationResponse, QuerySchemasResponse]>(
-        operations,
-      );
+      this.call<
+        [QueryServerInformationResponse, QuerySchemasResponse<TEntityTypeMap>]
+      >(operations);
 
     /**
      * Resolved once session is initialized.
@@ -244,7 +251,7 @@ export class Session {
 
     this.schemasPromise = initializingPromise
       .then((responses) => responses[1])
-      .catch(() => [] as Schema[]);
+      .catch(() => [] as Schema<TEntityTypeMap, keyof TEntityTypeMap>[]);
   }
 
   /**
@@ -252,7 +259,7 @@ export class Session {
    *
    * @return {Array|null} List of primary key attributes.
    */
-  getPrimaryKeyAttributes(entityType: string): string[] | null {
+  getPrimaryKeyAttributes(entityType: keyof TEntityTypeMap): string[] | null {
     // Todo: make this async in next major
     if (!this.schemas) {
       logger.warn("Schemas not available.");
@@ -549,7 +556,7 @@ export class Session {
   }
 
   /** Return encoded *operations*. */
-  encodeOperations(operations: operation.Operation[]) {
+  encodeOperations(operations: operation.Operation<keyof TEntityTypeMap>[]) {
     return JSON.stringify(this.encode(operations));
   }
 
@@ -587,12 +594,14 @@ export class Session {
    * This is cached after the first call, and assumes that the schemas will not change during the session.
    * @returns Promise with the API schemas for the session.
    */
-  async getSchemas(): Promise<Schema[]> {
+  async getSchemas(): Promise<Schema<TEntityTypeMap[keyof TEntityTypeMap]>[]> {
     if (!this.schemasPromise) {
-      this.schemasPromise = this.call<QuerySchemasResponse>([
+      this.schemasPromise = this.call<QuerySchemasResponse<TEntityTypeMap>>([
         { action: "query_schemas" },
       ]).then((responses) => {
-        this.schemaMapping = getSchemaMappingFromSchemas(responses[0]);
+        this.schemaMapping = getSchemaMappingFromSchemas<
+          TEntityTypeMap[keyof TEntityTypeMap]
+        >(responses[0]);
         return responses[0];
       });
     }
@@ -624,8 +633,8 @@ export class Session {
    * @param {string} options.decodeDatesAsIso - Return dates as ISO strings instead of moment objects
    *
    */
-  async call<T = ActionResponse>(
-    operations: operation.Operation[],
+  async call<T = ActionResponse<keyof TEntityTypeMap>>(
+    operations: operation.Operation<keyof TEntityTypeMap>[],
     {
       abortController,
       pushToken,
@@ -725,41 +734,45 @@ export class Session {
    *   Return update or create promise.
    */
 
-  ensure<T extends Data = Data>(
-    entityType: string,
-    data: T,
-    identifyingKeys: Array<keyof T> = [],
-  ): Promise<T & Entity> {
+  ensure<TEntityType extends keyof TEntityTypeMap>(
+    entityType: TEntityType,
+    data: TEntityTypeMap[TEntityType],
+    identifyingKeys: Array<keyof TEntityTypeMap[TEntityType]> = [],
+  ): Promise<TEntityTypeMap[TEntityType]> {
     let keys = identifyingKeys as string[];
 
+    const anyData = data as any;
+
     logger.info(
       "Ensuring entity with data using identifying keys: ",
       entityType,
-      data,
+      anyData,
       keys,
     );
 
     if (!keys.length) {
-      keys = Object.keys(data);
+      keys = Object.keys(anyData);
     }
 
     if (!keys.length) {
       throw new Error(
         "Could not determine any identifying data to check against " +
-          `when ensuring ${entityType} with data ${data}. ` +
+          `when ensuring ${String(entityType)} with data ${anyData}. ` +
           `Identifying keys: ${identifyingKeys}`,
       );
     }
 
     const primaryKeys = this.getPrimaryKeyAttributes(entityType);
     if (primaryKeys === null || primaryKeys.length === 0) {
-      throw new Error(`Primary key could not be found for: ${entityType}`);
+      throw new Error(
+        `Primary key could not be found for: ${String(entityType)}`,
+      );
     }
     let expression = `select ${primaryKeys.join(
       ", ",
-    )} from ${entityType} where`;
+    )} from ${String(entityType)} where`;
     const criteria = keys.map((identifyingKey) => {
-      let value = data[identifyingKey];
+      let value = anyData[identifyingKey];
 
       if (value != null && typeof value.valueOf() === "string") {
         value = `"${value}"`;
@@ -772,48 +785,51 @@ export class Session {
 
     expression = `${expression} ${criteria.join(" and ")}`;
 
-    return this.query<T & Entity>(expression).then((response) => {
+    return this.query<TEntityType>(expression).then((response) => {
       if (response.data.length === 0) {
-        return this.create<T>(entityType, data).then(({ data: responseData }) =>
-          Promise.resolve(responseData),
+        return this.create<TEntityType>(entityType, anyData).then(
+          ({ data: responseData }) => Promise.resolve(responseData),
         );
       }
 
       if (response.data.length !== 1) {
         throw new Error(
           "Expected single or no item to be found but got multiple " +
-            `when ensuring ${entityType} with data ${data}. ` +
+            `when ensuring ${String(entityType)} with data ${anyData}. ` +
             `Identifying keys: ${identifyingKeys}`,
         );
       }
 
-      const updateEntity = response.data[0];
+      const updateEntity = response.data[0] as any;
 
       // Update entity if required.
       let updated = false;
-      Object.keys(data).forEach((key: keyof (T & Entity)) => {
-        if (data[key] !== updateEntity[key]) {
-          updateEntity[key] = data[key];
+      Object.keys(anyData).forEach((key) => {
+        if (anyData[key] !== updateEntity[key]) {
+          updateEntity[key] = anyData[key];
           updated = true;
         }
       });
 
       if (updated) {
-        return this.update<T>(
+        return this.update<TEntityType>(
           entityType,
           primaryKeys.map((key: string) => updateEntity[key]),
-          Object.keys(data).reduce<T>((accumulator, key: keyof T) => {
-            if (primaryKeys.indexOf(key.toString()) === -1) {
-              accumulator[key] = data[key];
-            }
-            return accumulator;
-          }, {} as T),
+          Object.keys(anyData).reduce<any>(
+            (accumulator, key) => {
+              if (primaryKeys.indexOf(key.toString()) === -1) {
+                accumulator[key] = anyData[key];
+              }
+              return accumulator;
+            },
+            {} as TEntityTypeMap[TEntityType],
+          ),
         ).then(({ data: responseData }) =>
-          Promise.resolve(responseData as T & Entity),
+          Promise.resolve(responseData as TEntityTypeMap[TEntityType]),
         );
       }
 
-      return Promise.resolve(response.data[0]);
+      return Promise.resolve(response.data[0]) as any;
     });
   }
 
@@ -822,7 +838,9 @@ export class Session {
    * @param  {string} schemaId Id of schema model, e.g. `AssetVersion`.
    * @return {Object|null} Schema definition
    */
-  getSchema(schemaId: string): Schema | null {
+  getSchema<TEntityType extends keyof TEntityTypeMap>(
+    schemaId: TEntityType,
+  ): Schema<TEntityTypeMap, TEntityType> | null {
     // TODO: make this async in next major
     const schema = this.schemaMapping?.[schemaId];
     return schema ?? null;
@@ -841,15 +859,14 @@ export class Session {
    * @return {Promise} Promise which will be resolved with an object
    * containing action, data and metadata
    */
-  async query<T extends Data = Data>(
+  async query<TEntityType extends keyof TEntityTypeMap>(
     expression: string,
     options: QueryOptions = {},
   ) {
     logger.debug("Query", expression);
-    const responses = await this.call<[QueryResponse<T>]>(
-      [operation.query(expression)],
-      options,
-    );
+    const responses = await this.call<
+      [QueryResponse<TEntityTypeMap[TEntityType]>]
+    >([operation.query(expression)], options);
     return responses[0];
   }
 
@@ -871,14 +888,14 @@ export class Session {
    * @return {Promise} Promise which will be resolved with an object
    * containing data and metadata
    */
-  async search<T extends Data = Data>(
+  async search<TEntityType extends keyof TEntityTypeMap>(
     {
       expression,
       entityType,
       terms = [],
       contextId,
       objectTypeIds,
-    }: SearchOptions,
+    }: SearchOptions<TEntityType>,
     options: QueryOptions = {},
   ) {
     logger.debug("Search", {
@@ -889,7 +906,7 @@ export class Session {
       objectTypeIds,
     });
 
-    const responses = await this.call<[SearchResponse<T>]>(
+    const responses = await this.call<[SearchResponse<TEntityType>]>(
       [
         operation.search({
           expression,
@@ -916,17 +933,16 @@ export class Session {
    * @param {object} options.ensureSerializableResponse - Disable normalization of response data
    * @return {Promise} Promise which will be resolved with the response.
    */
-  async create<T extends Data = Data>(
-    entityType: string,
-    data: T,
+  async create<TEntityType extends keyof TEntityTypeMap>(
+    entityType: TEntityType,
+    data: Partial<TEntityTypeMap[TEntityType]>,
     options: MutationOptions = {},
   ) {
     logger.debug("Create", entityType, data, options);
 
-    const responses = await this.call<[CreateResponse<T & Entity>]>(
-      [operation.create(entityType, data)],
-      options,
-    );
+    const responses = await this.call<
+      [CreateResponse<TEntityTypeMap[TEntityType]>]
+    >([operation.create(entityType, data)], options);
     return responses[0];
   }
 
@@ -943,18 +959,17 @@ export class Session {
    * @param {object} options.ensureSerializableResponse - Disable normalization of response data
    * @return {Promise} Promise resolved with the response.
    */
-  async update<T extends Data = Data>(
-    type: string,
+  async update<TEntityType extends keyof TEntityTypeMap>(
+    type: TEntityType,
     keys: string[] | string,
-    data: T,
+    data: Partial<TEntityTypeMap[TEntityType]>,
     options: MutationOptions = {},
   ) {
     logger.debug("Update", type, keys, data, options);
 
-    const responses = await this.call<[UpdateResponse<T>]>(
-      [operation.update(type, keys, data)],
-      options,
-    );
+    const responses = await this.call<
+      [UpdateResponse<TEntityTypeMap[TEntityType]>]
+    >([operation.update(type, keys, data)], options);
     return responses[0];
   }
 
@@ -969,8 +984,8 @@ export class Session {
    * @param {object} options.decodeDatesAsIso - Decode dates as ISO strings instead of moment objects
    * @return {Promise} Promise resolved with the response.
    */
-  async delete(
-    type: string,
+  async delete<TEntityType extends keyof TEntityTypeMap>(
+    type: TEntityType,
     keys: string[] | string,
     options: MutationOptions = {},
   ) {
@@ -1053,10 +1068,14 @@ export class Session {
     file: Blob,
     options: CreateComponentOptions = {},
   ): Promise<
-    readonly [CreateResponse, CreateResponse, GetUploadMetadataResponse]
+    readonly [
+      CreateResponse<TEntityTypeMap["FileComponent"]>,
+      CreateResponse<TEntityTypeMap["ComponentLocation"]>,
+      GetUploadMetadataResponse,
+    ]
   > {
     return new Promise((resolve, reject) => {
-      const uploader = new Uploader(this, file, {
+      const uploader = new Uploader<TEntityTypeMap>(this, file, {
         ...options,
         onError(error) {
           reject(error);
diff --git a/source/types.ts b/source/types.ts
index d880b14..b986634 100644
--- a/source/types.ts
+++ b/source/types.ts
@@ -25,43 +25,39 @@ export interface CreateComponentOptions {
   onAborted?: () => unknown;
 }
 
-export interface Entity {
-  id: string;
-  __entity_type__: string;
-}
 interface ResponseMetadata {
   next: {
     offset: number | null;
   };
 }
-export interface SearchOptions {
+export interface SearchOptions<TEntityType> {
   expression: string;
-  entityType: string;
+  entityType: TEntityType;
   terms?: string[];
   contextId?: string;
   objectTypeIds?: string[];
 }
 
-export interface QueryResponse<T = Data> {
-  data: T[];
+export interface QueryResponse<TEntityData> {
+  data: TEntityData[];
   action: "query";
   metadata: ResponseMetadata;
 }
 
-export interface CreateResponse<T = Data> {
-  data: T;
+export interface CreateResponse<TEntityData> {
+  data: TEntityData;
   action: "create";
 }
-export interface UpdateResponse<T = Data> {
-  data: T;
+export interface UpdateResponse<TEntityData> {
+  data: TEntityData;
   action: "update";
 }
 export interface DeleteResponse {
   data: true;
   action: "delete";
 }
-export interface SearchResponse<T = Data> {
-  data: T[];
+export interface SearchResponse<TEntityData> {
+  data: TEntityData[];
   action: "search";
   metadata: ResponseMetadata;
 }
@@ -69,7 +65,7 @@ export interface ResetRemoteResponse {
   action: "reset_remote";
   data: Data;
 }
-export type QuerySchemasResponse = Schema[];
+export type QuerySchemasResponse<TEntityTypeMap> = Schema<TEntityTypeMap>[];
 
 export type QueryServerInformationResponse = ServerInformation;
 export interface ServerInformation {
@@ -149,14 +145,18 @@ export interface PermissionsResponse {
   data: Data;
 }
 
-export type ActionResponse =
-  | QueryResponse
-  | CreateResponse
-  | UpdateResponse
+export type ActionResponse<
+  TEntityTypeMap,
+  TEntityType extends keyof TEntityTypeMap = keyof TEntityTypeMap,
+  TEntityData = TEntityTypeMap[TEntityType],
+> =
+  | QueryResponse<TEntityData>
+  | CreateResponse<TEntityData>
+  | UpdateResponse<TEntityData>
   | DeleteResponse
-  | SearchResponse
+  | SearchResponse<TEntityData>
   | ResetRemoteResponse
-  | QuerySchemasResponse
+  | QuerySchemasResponse<TEntityTypeMap>
   | QueryServerInformationResponse
   | GetWidgetUrlResponse
   | DelayedJobResponse
@@ -201,21 +201,24 @@ export type TypedSchemaProperty =
 export type RefSchemaProperty = {
   ["$ref"]: string;
 };
-export type SchemaProperties = {
-  [key: string]: TypedSchemaProperty | RefSchemaProperty;
+export type SchemaProperties<TEntityData> = {
+  [key in keyof TEntityData]: TypedSchemaProperty | RefSchemaProperty;
 };
 export type SchemaMixin = {
   $ref: string;
 };
 export type SchemaMetadata = { entity_event: boolean };
-export interface Schema {
-  properties: SchemaProperties;
+export interface Schema<
+  TEntityTypeMap,
+  TEntityType extends keyof TEntityTypeMap = keyof TEntityTypeMap,
+> {
+  properties: SchemaProperties<TEntityTypeMap[TEntityType]>;
   default_projections: string[];
   primary_key: string[];
   required: string[];
   immutable: string[];
   type?: string;
-  id: string;
+  id: TEntityType;
   computed?: string[];
   system_projections?: string[];
   alias_for?: string | Data;
@@ -231,3 +234,10 @@ export interface QueryOptions {
 }
 
 export interface CallOptions extends MutationOptions, QueryOptions {}
+
+export interface ExtendibleEntityTypeMap {}
+export interface DefaultEntityTypeMap {
+  [key: string]: {
+    [key: string]: any;
+  };
+}
diff --git a/source/uploader.ts b/source/uploader.ts
index ba54470..2e16c42 100644
--- a/source/uploader.ts
+++ b/source/uploader.ts
@@ -14,7 +14,6 @@ import type {
 } from "./types.js";
 import normalizeString from "./util/normalize_string.js";
 import { splitFileExtension } from "./util/split_file_extension.js";
-import type { Data } from "./types.js";
 import { getChunkSize } from "./util/get_chunk_size.js";
 import { backOff } from "./util/back_off.js";
 
@@ -50,11 +49,11 @@ declare global {
  * 3. Completion - Multi-part uploads are completed, the component is
  *    registered in the server location and an event is published.
  */
-export class Uploader {
+export class Uploader<TEntityTypeMap extends Record<string, any>> {
   /** Component id */
   componentId: string;
   /** Session instance */
-  private session: Session;
+  private session: Session<TEntityTypeMap>;
   /** File to upload */
   private file: Blob;
   /** Called on upload progress with percentage */
@@ -98,13 +97,21 @@ export class Uploader {
   /** Additional data for Component entity */
   private data: CreateComponentOptions["data"];
   /** @deprecated - Remove once Session.createComponent signature is updated. */
-  createComponentResponse: CreateResponse<Data> | null;
+  createComponentResponse: CreateResponse<
+    TEntityTypeMap["FileComponent"]
+  > | null;
   /** @deprecated - Remove once Session.createComponent signature is updated. */
   uploadMetadata: GetUploadMetadataResponse | null;
   /** @deprecated - Remove once Session.createComponent signature is updated. */
-  createComponentLocationResponse: CreateResponse<Data> | null;
+  createComponentLocationResponse: CreateResponse<
+    TEntityTypeMap["ComponentLocation"]
+  > | null;
 
-  constructor(session: Session, file: Blob, options: UploaderOptions) {
+  constructor(
+    session: Session<TEntityTypeMap>,
+    file: Blob,
+    options: UploaderOptions,
+  ) {
     this.session = session;
     this.file = file;
     const componentName = options.name ?? (file as File).name;
@@ -196,7 +203,10 @@ export class Uploader {
       size: this.fileSize,
     };
     const response = await this.session.call<
-      [CreateResponse, GetUploadMetadataResponse]
+      [
+        CreateResponse<TEntityTypeMap["FileComponent"]>,
+        GetUploadMetadataResponse,
+      ]
     >([
       operation.create("FileComponent", component),
       {
@@ -531,7 +541,10 @@ export class Uploader {
       }),
     );
 
-    const response = await this.session.call<CreateResponse>(operations);
+    const response =
+      await this.session.call<
+        CreateResponse<TEntityTypeMap["ComponentLocation"]>
+      >(operations);
     this.createComponentLocationResponse = response[response.length - 1];
 
     // Emit event so that clients can perform additional work on uploaded
diff --git a/source/util/get_schema_mapping.ts b/source/util/get_schema_mapping.ts
index 4dea004..b7d8ffc 100644
--- a/source/util/get_schema_mapping.ts
+++ b/source/util/get_schema_mapping.ts
@@ -2,8 +2,14 @@
 
 import type { Schema } from "../types.js";
 
-export default function getSchemaMappingFromSchemas(schemas: Schema[]) {
-  const schemaMapping = {} as Record<string, Schema>;
+export default function getSchemaMappingFromSchemas<TEntityTypeMap>(
+  schemas: Schema<TEntityTypeMap>[],
+): {
+  [TEntityType in keyof TEntityTypeMap]: Schema<TEntityTypeMap, TEntityType>;
+} {
+  const schemaMapping = {} as {
+    [TEntityType in keyof TEntityTypeMap]: Schema<TEntityTypeMap, TEntityType>;
+  };
   for (const schema of schemas) {
     schemaMapping[schema.id] = schema;
   }
diff --git a/test/session.test.ts b/test/session.test.ts
index 6f2849f..b2427b7 100755
--- a/test/session.test.ts
+++ b/test/session.test.ts
@@ -290,7 +290,9 @@ describe("Session", () => {
   });
 
   it("Should decode batched query operations", async () => {
-    const responses = await session.call<QueryResponse>([
+    const responses = await session.call<
+      QueryResponse<{ status: { state: { short: string } } }>
+    >([
       operation.query(
         "select status.state.short from Task where status.state.short is NOT_STARTED limit 1",
       ),