captify 0.1.18 → 0.1.19

package/dist/llm-providers/services/llm-providers.service.js

@@ -1,17 +1,19 @@
-import { createAmazonBedrock } from "@ai-sdk/amazon-bedrock";
-import { createAnthropic } from "@ai-sdk/anthropic";
-import { createOpenAI } from "@ai-sdk/openai";
-import { convertToModelMessages, extractReasoningMiddleware, generateText, streamText, wrapLanguageModel, } from "ai";
-import { sanitizeMessagesForBedrock } from "../../shared/utilities/file";
-import { TABLES } from "../../shared/shared.constants";
-import { DynamoDbInstances } from "../../shared/services/dynamodb";
-import { UpdateExpressionBuilder } from "../../shared/services/dynamodb";
-import { buildModelSK, buildProviderSK, buildTenantPK, LLM_KEYS, LLM_ERROR_MESSAGES, } from "../llm-providers.constants";
+import { llmProviderKeys } from "../../db/schemas";
+import { LLM_ERROR_MESSAGES, LLM_KEYS, THINKING_BUDGET, } from "../llm-providers.constants";
 import { InvalidProviderTypeError, ModelNotAvailableError, ModelNotFoundError, ProviderNotFoundError, } from "../llm-providers.errors";
 import { createModelParamsSchema, createProviderParamsSchema, deleteModelParamsSchema, deleteProviderParamsSchema, getModelParamsSchema, getProviderByIdParamsSchema, llmRequestParamsSchema, modelSchema, providerSchema, updateModelParamsSchema, updateProviderParamsSchema, } from "../llm-providers.schemas";
-import { LoggerFactory } from "../../shared/services/logger";
+import { getUsageTrackingService } from "../../model-budget/services";
 import { ValidationError } from "../../shared/services";
 import { BaseAwsService } from "../../shared/services/base-aws-service";
+import { DynamoDbInstances } from "../../shared/services/dynamodb";
+import { LoggerFactory } from "../../shared/services/logger";
+import { TABLES, USER_CONTEXT_WITH_METADATA } from "../../shared/shared.constants";
+import { sanitizeMessagesForBedrock } from "../../shared/utilities/file";
+import { normalizeToolParts, sanitizeToolPartsForBedrock, stripReasoningFromMessages, } from "../../shared/utilities/messages";
+import { createAmazonBedrock } from "@ai-sdk/amazon-bedrock";
+import { createAnthropic } from "@ai-sdk/anthropic";
+import { createOpenAI } from "@ai-sdk/openai";
+import { convertToModelMessages, generateText, streamText } from "ai";
 /**
  * Service for managing LLM providers and models.
  * Handles AI model interactions through the Vercel AI SDK.
@@ -32,20 +34,21 @@ export class LLMService extends BaseAwsService {
      */
     async getAllProviders(tenantId, options) {
         this.logger.debug("Fetching all providers", { tenantId });
-        const { Items } = await this.db.queryItems({
-            KeyConditionExpression: "PK = :pk AND begins_with(SK, :prefix)",
-            ExpressionAttributeValues: {
-                ":pk": buildTenantPK(tenantId),
-                ":prefix": LLM_KEYS.SORT_KEY.METADATA,
+        const { items: Items } = await this.db.queryTyped({
+            keyCondition: {
+                PK: llmProviderKeys.pk(tenantId),
+                SK: { beginsWith: LLM_KEYS.SORT_KEY.METADATA },
             },
         }, { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
-        return (Items !== null && Items !== void 0 ? Items : []).filter((item) => {
-            const result = providerSchema.safeParse(item);
+        return (Items !== null && Items !== void 0 ? Items : [])
+            .map((item) => providerSchema.safeParse(item))
+            .filter((result) => {
             if (!result.success) {
-                this.logger.warn(`Invalid provider item found in getAllProviders: ${JSON.stringify(item)}`);
+                this.logger.warn(`Invalid provider item found in getAllProviders: ${JSON.stringify(result)}`);
             }
             return result.success;
-        });
+        })
+            .map((result) => result.data);
     }
     /**
      * Get all available providers that are enabled.
@@ -79,8 +82,8 @@ export class LLMService extends BaseAwsService {
         this.logger.debug("Fetching provider by ID", { tenantId, providerId });
         const validated = getProviderByIdParamsSchema.parse({ tenantId, providerId });
         const { Item } = await this.db.getItem({
-            PK: buildTenantPK(validated.tenantId),
-            SK: buildProviderSK(validated.providerId),
+            PK: llmProviderKeys.pk(validated.tenantId),
+            SK: llmProviderKeys.providerSk(validated.providerId),
         }, { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
         if (!Item)
             return null;
@@ -111,6 +114,7 @@ export class LLMService extends BaseAwsService {
         const now = Date.now();
         const newProvider = {
             id,
+            createdById: validated.provider.createdById,
             name: validated.provider.name,
             type: validated.provider.type,
             description: validated.provider.description,
@@ -119,7 +123,7 @@ export class LLMService extends BaseAwsService {
             createdAt: now,
             updatedAt: now,
         };
-        await this.db.putItem(Object.assign({ PK: buildTenantPK(validated.tenantId), SK: buildProviderSK(id) }, newProvider), { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
+        await this.db.putItem(Object.assign({ PK: llmProviderKeys.pk(validated.tenantId), SK: llmProviderKeys.providerSk(id), itemType: "PROVIDER" }, newProvider), { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
         this.logger.info("Created provider", {
             id: newProvider.id,
             name: newProvider.name,
@@ -152,8 +156,10 @@ export class LLMService extends BaseAwsService {
         if (!current) {
             throw new ProviderNotFoundError(validated.providerId);
         }
-        const builder = new UpdateExpressionBuilder()
+        const builder = this.db
+            .createUpdateBuilder()
             .set("updatedAt", Date.now())
+            .setOptional("modifiedById", validated.updates.modifiedById)
             .setOptional("name", validated.updates.name)
             .setOptional("type", validated.updates.type)
             .setOptional("description", validated.updates.description)
@@ -161,8 +167,8 @@ export class LLMService extends BaseAwsService {
             .setOptional("isEnabled", validated.updates.isEnabled);
         const { UpdateExpression, ExpressionAttributeNames, ExpressionAttributeValues } = builder.build();
         const { Attributes } = await this.db.updateItem({
-            PK: buildTenantPK(validated.tenantId),
-            SK: buildProviderSK(validated.providerId),
+            PK: llmProviderKeys.pk(validated.tenantId),
+            SK: llmProviderKeys.providerSk(validated.providerId),
         }, UpdateExpression, ExpressionAttributeNames, ExpressionAttributeValues, { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
         const result = providerSchema.safeParse(Attributes);
         if (!result.success) {
@@ -202,8 +208,8 @@ export class LLMService extends BaseAwsService {
             throw new ValidationError(`Cannot delete provider ${provider.name}: it has ${models.length} model(s)`);
         }
         await this.db.deleteItem({
-            PK: buildTenantPK(validated.tenantId),
-            SK: buildProviderSK(validated.providerId),
+            PK: llmProviderKeys.pk(validated.tenantId),
+            SK: llmProviderKeys.providerSk(validated.providerId),
         }, { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
         this.logger.info("Deleted provider", { id: provider.id, name: provider.name });
     }
@@ -222,14 +228,16 @@ export class LLMService extends BaseAwsService {
      */
     async getAllModels(tenantId, options) {
         this.logger.debug("Fetching all models", { tenantId });
-        const { Items } = await this.db.queryItems({
-            KeyConditionExpression: "PK = :pk AND begins_with(SK, :prefix)",
-            ExpressionAttributeValues: {
-                ":pk": buildTenantPK(tenantId),
-                ":prefix": LLM_KEYS.SORT_KEY.MODEL_PREFIX,
+        const { items: Items } = await this.db.queryTyped({
+            keyCondition: {
+                PK: llmProviderKeys.pk(tenantId),
+                SK: { beginsWith: LLM_KEYS.SORT_KEY.MODEL_PREFIX },
             },
         }, { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
-        return (Items !== null && Items !== void 0 ? Items : []).filter((item) => modelSchema.safeParse(item).success);
+        return (Items !== null && Items !== void 0 ? Items : [])
+            .map((item) => modelSchema.safeParse(item))
+            .filter((result) => result.success)
+            .map((result) => result.data);
     }
     /**
      * Get all providers with their models in a single efficient query.
@@ -253,10 +261,9 @@ export class LLMService extends BaseAwsService {
     async getProvidersWithModels(tenantId, enabledOnly = true, availableOnly = true, options) {
         this.logger.debug("Fetching providers with models", { tenantId, enabledOnly, availableOnly });
         // Single query to get all providers and models for the tenant (most efficient)
-        const { Items } = await this.db.queryItems({
-            KeyConditionExpression: "PK = :pk",
-            ExpressionAttributeValues: {
-                ":pk": buildTenantPK(tenantId),
+        const { items: Items } = await this.db.queryTyped({
+            keyCondition: {
+                PK: llmProviderKeys.pk(tenantId),
             },
         }, { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
         const providers = [];
@@ -293,11 +300,11 @@ export class LLMService extends BaseAwsService {
         });
     }
     /**
-     * Get a specific model by provider and model ID.
+     * Get a specific model by provider and model record ID (UUID).
      *
      * @param tenantId - Tenant ID to scope the session
      * @param providerId Provider UUID
-     * @param modelId Model UUID
+     * @param modelId Model record UUID (not the actual model identifier like "gpt-4")
      * @param options Optional service options including user credentials
      * @returns Promise resolving to model or null if not found
      * @throws ValidationError if providerId or modelId is invalid
@@ -306,7 +313,7 @@ export class LLMService extends BaseAwsService {
      * ```typescript
      * const model = await llmService.getModel(tenantId, providerUuid, modelUuid, { userCredentials });
      * if (model) {
-     *   console.log(model.name, model.capabilities);
+     *   console.log(model.name, model.modelId, model.capabilities);
      * }
      * ```
      */
@@ -314,14 +321,29 @@ export class LLMService extends BaseAwsService {
         this.logger.debug("Fetching model", { tenantId, providerId, modelId });
         const validated = getModelParamsSchema.parse({ tenantId, providerId, modelId });
         const { Item } = await this.db.getItem({
-            PK: buildTenantPK(validated.tenantId),
-            SK: buildModelSK(validated.providerId, validated.modelId),
+            PK: llmProviderKeys.pk(validated.tenantId),
+            SK: llmProviderKeys.modelSk(validated.providerId, validated.modelId),
         }, { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
         if (!Item)
             return null;
         const result = modelSchema.safeParse(Item);
         return result.success ? result.data : null;
     }
+    /**
+     * Look up a model by its UUID (record ID) without requiring a provider ID.
+     * Scans all tenant models and returns the match. Suitable for collaboration
+     * flows where only the model UUID is sent from the client.
+     *
+     * @param tenantId - Tenant ID to scope the query
+     * @param modelId - Model record UUID
+     * @param options Optional service options including user credentials
+     * @returns Promise resolving to model or null if not found
+     */
+    async getModelById(tenantId, modelId, options) {
+        var _a;
+        const models = await this.getAllModels(tenantId, options);
+        return (_a = models.find((m) => m.id === modelId)) !== null && _a !== void 0 ? _a : null;
+    }
     /**
      * Get the count of models for a specific provider.
      *
@@ -340,15 +362,16 @@ export class LLMService extends BaseAwsService {
     async getModelCountByProvider(tenantId, providerId, options) {
         this.logger.debug("Fetching model count by provider", { tenantId, providerId });
         const validated = getProviderByIdParamsSchema.parse({ tenantId, providerId });
-        const { Count } = await this.db.queryItems({
-            KeyConditionExpression: "PK = :pk AND begins_with(SK, :prefix)",
-            ExpressionAttributeValues: {
-                ":pk": buildTenantPK(validated.tenantId),
-                ":prefix": `${LLM_KEYS.SORT_KEY.MODEL_PREFIX}${validated.providerId}${LLM_KEYS.SEPARATOR}`,
+        const { count } = await this.db.queryTyped({
+            keyCondition: {
+                PK: llmProviderKeys.pk(validated.tenantId),
+                SK: {
+                    beginsWith: `${LLM_KEYS.SORT_KEY.MODEL_PREFIX}${validated.providerId}${LLM_KEYS.SEPARATOR}`,
+                },
             },
-            Select: "COUNT",
+            select: "COUNT",
         }, { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
-        return Count !== null && Count !== void 0 ? Count : 0;
+        return count !== null && count !== void 0 ? count : 0;
     }
     /**
      * Get all models for a specific provider.
@@ -362,14 +385,18 @@ export class LLMService extends BaseAwsService {
     async getModelsByProvider(tenantId, providerId, options) {
         this.logger.debug("Fetching models by provider", { tenantId, providerId });
         const validated = getProviderByIdParamsSchema.parse({ tenantId, providerId });
-        const { Items } = await this.db.queryItems({
-            KeyConditionExpression: "PK = :pk AND begins_with(SK, :prefix)",
-            ExpressionAttributeValues: {
-                ":pk": buildTenantPK(validated.tenantId),
-                ":prefix": `${LLM_KEYS.SORT_KEY.MODEL_PREFIX}${validated.providerId}${LLM_KEYS.SEPARATOR}`,
+        const { items: Items } = await this.db.queryTyped({
+            keyCondition: {
+                PK: llmProviderKeys.pk(validated.tenantId),
+                SK: {
+                    beginsWith: `${LLM_KEYS.SORT_KEY.MODEL_PREFIX}${validated.providerId}${LLM_KEYS.SEPARATOR}`,
+                },
             },
         }, { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
-        return (Items !== null && Items !== void 0 ? Items : []).filter((item) => modelSchema.safeParse(item).success);
+        return (Items !== null && Items !== void 0 ? Items : [])
+            .map((item) => modelSchema.safeParse(item))
+            .filter((result) => result.success)
+            .map((result) => result.data);
     }
     /**
      * Create a new model (admin only).
@@ -399,23 +426,29 @@ export class LLMService extends BaseAwsService {
         if (!provider) {
             throw new ProviderNotFoundError(validated.model.providerId);
         }
-        const existingModel = await this.getModel(validated.tenantId, validated.model.providerId, validated.model.id, options);
-        if (existingModel) {
-            throw new ValidationError(`Model with ID "${validated.model.id}" already exists for this provider`);
+        // Check for duplicate modelId within this provider
+        const existingModels = await this.getModelsByProvider(validated.tenantId, validated.model.providerId, options);
+        const duplicateModel = existingModels.find((m) => m.modelId === validated.model.modelId);
+        if (duplicateModel) {
+            throw new ValidationError(`Model with identifier "${validated.model.modelId}" already exists for this provider`);
         }
+        const id = crypto.randomUUID();
         const now = Date.now();
         // If setting as Basic Assistant default, clear any existing default first
         if (validated.model.isBasicAssistantDefault) {
             await this.clearBasicAssistantDefault(validated.tenantId, options);
         }
         const newModel = {
-            id: validated.model.id,
+            id,
+            createdById: validated.model.createdById,
+            modelId: validated.model.modelId,
             providerId: validated.model.providerId,
             name: validated.model.name,
             description: validated.model.description,
             capabilities: validated.model.capabilities,
             contextWindow: validated.model.contextWindow,
             maxOutputTokens: validated.model.maxOutputTokens,
+            thinkingBudgetTokens: validated.model.thinkingBudgetTokens,
             pricing: validated.model.pricing,
             isAvailable: validated.model.isAvailable,
             isBasicAssistantDefault: (_a = validated.model.isBasicAssistantDefault) !== null && _a !== void 0 ? _a : false,
@@ -423,9 +456,10 @@ export class LLMService extends BaseAwsService {
             createdAt: now,
             updatedAt: now,
         };
-        await this.db.putItem(Object.assign({ PK: buildTenantPK(validated.tenantId), SK: buildModelSK(validated.model.providerId, validated.model.id) }, newModel), { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
+        await this.db.putItem(Object.assign({ PK: llmProviderKeys.pk(validated.tenantId), SK: llmProviderKeys.modelSk(validated.model.providerId, id), itemType: "MODEL" }, newModel), { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
         this.logger.info("Created model", {
             id: newModel.id,
+            modelId: newModel.modelId,
             name: newModel.name,
             providerId: newModel.providerId,
             isBasicAssistantDefault: newModel.isBasicAssistantDefault,
@@ -437,19 +471,18 @@ export class LLMService extends BaseAwsService {
      *
      * @param tenantId - Tenant ID to scope the session
      * @param providerId Provider UUID
-     * @param modelId Model UUID
-     * @param updates Model data to update
+     * @param modelId Model record UUID (not the actual model identifier)
+     * @param updates Model data to update (including modelId for actual model identifier)
      * @param options Optional service options including user credentials
      * @returns Promise resolving to saved model
-     * @throws ValidationError if model data is invalid
+     * @throws ValidationError if model data is invalid or new modelId already exists
      *
      * @example
      * ```typescript
      * const model = await llmService.updateModel(tenantId, providerUuid, modelUuid, {
-     *   name: "GPT-4o",
+     *   modelId: "gpt-4-turbo-2024-04-09", // Update the actual model identifier
+     *   name: "GPT-4 Turbo",
      *   capabilities: ["chat", "vision"],
-     *   contextWindow: 128000,
-     *   maxOutputTokens: 4096,
      *   isAvailable: true,
      * });
      * ```
@@ -460,13 +493,24 @@ export class LLMService extends BaseAwsService {
         if (!current) {
             throw new ModelNotFoundError(validated.modelId);
         }
+        // If updating modelId (actual model identifier), check for duplicates
+        if (validated.updates.modelId && validated.updates.modelId !== current.modelId) {
+            const existingModels = await this.getModelsByProvider(validated.tenantId, validated.providerId, options);
+            const duplicateModel = existingModels.find((m) => m.modelId === validated.updates.modelId && m.id !== current.id);
+            if (duplicateModel) {
+                throw new ValidationError(`Model with identifier "${validated.updates.modelId}" already exists for this provider`);
+            }
+        }
         // If setting as Basic Assistant default, clear any existing default first
         // (but don't clear if it's already the default and we're just updating other fields)
         if (validated.updates.isBasicAssistantDefault === true && !current.isBasicAssistantDefault) {
             await this.clearBasicAssistantDefault(validated.tenantId, options);
         }
-        const builder = new UpdateExpressionBuilder()
+        const builder = this.db
+            .createUpdateBuilder()
             .set("updatedAt", Date.now())
+            .setOptional("modifiedById", validated.updates.modifiedById)
+            .setOptional("modelId", validated.updates.modelId)
             .setOptional("name", validated.updates.name)
             .setOptional("description", validated.updates.description)
             .setOptional("capabilities", validated.updates.capabilities)
@@ -476,10 +520,24 @@ export class LLMService extends BaseAwsService {
             .setOptional("parameterConstraints", validated.updates.parameterConstraints)
             .setOptional("isAvailable", validated.updates.isAvailable)
             .setOptional("isBasicAssistantDefault", validated.updates.isBasicAssistantDefault);
+        // Handle thinkingBudgetTokens specially: null means remove, undefined means skip, value means set
+        if (validated.updates.thinkingBudgetTokens === null) {
+            builder.remove("thinkingBudgetTokens");
+        }
+        else if (validated.updates.thinkingBudgetTokens !== undefined) {
+            builder.set("thinkingBudgetTokens", validated.updates.thinkingBudgetTokens);
+        }
+        // Handle budgetConfig specially: null means remove, undefined means skip, value means set
+        if (validated.updates.budgetConfig === null) {
+            builder.remove("budgetConfig");
+        }
+        else if (validated.updates.budgetConfig !== undefined) {
+            builder.set("budgetConfig", validated.updates.budgetConfig);
+        }
         const { UpdateExpression, ExpressionAttributeNames, ExpressionAttributeValues } = builder.build();
         const { Attributes } = await this.db.updateItem({
-            PK: buildTenantPK(validated.tenantId),
-            SK: buildModelSK(validated.providerId, validated.modelId),
+            PK: llmProviderKeys.pk(validated.tenantId),
+            SK: llmProviderKeys.modelSk(validated.providerId, validated.modelId),
         }, UpdateExpression, ExpressionAttributeNames, ExpressionAttributeValues, { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
         const result = modelSchema.safeParse(Attributes);
         if (!result.success) {
@@ -487,6 +545,7 @@ export class LLMService extends BaseAwsService {
         }
         this.logger.info("Updated model", {
             id: result.data.id,
+            modelId: result.data.modelId,
             name: result.data.name,
             providerId: result.data.providerId,
             isBasicAssistantDefault: result.data.isBasicAssistantDefault,
@@ -498,7 +557,7 @@ export class LLMService extends BaseAwsService {
      *
      * @param tenantId - Tenant ID to scope the session
      * @param providerId Provider UUID
-     * @param modelId Model UUID to delete
+     * @param modelId Model record UUID to delete
      * @param options Optional service options including user credentials
      * @returns Promise that resolves when deletion is complete
      * @throws ValidationError if providerId or modelId is invalid
@@ -509,9 +568,9 @@ export class LLMService extends BaseAwsService {
      *
      * @example
      * ```typescript
-     * await llmService.deleteModel(tenantId, providerId, modelId);
+     * await llmService.deleteModel(tenantId, providerId, modelUuid);
      * // Calling again won't throw an error
-     * await llmService.deleteModel(tenantId, providerId, modelId); // Silent success
+     * await llmService.deleteModel(tenantId, providerId, modelUuid); // Silent success
      * ```
      */
     async deleteModel(tenantId, providerId, modelId, options) {
@@ -521,11 +580,12 @@ export class LLMService extends BaseAwsService {
             return;
         }
         await this.db.deleteItem({
-            PK: buildTenantPK(validated.tenantId),
-            SK: buildModelSK(validated.providerId, validated.modelId),
+            PK: llmProviderKeys.pk(validated.tenantId),
+            SK: llmProviderKeys.modelSk(validated.providerId, validated.modelId),
         }, { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
         this.logger.info("Deleted model", {
             id: model.id,
+            modelId: model.modelId,
             name: model.name,
             providerId: model.providerId,
         });
@@ -585,6 +645,7 @@ export class LLMService extends BaseAwsService {
      * ```
      */
     async streamResponse(tenantId, providerId, modelId, messages, systemPrompt = "You are a helpful assistant.", options) {
+        var _a, _b;
         const validated = llmRequestParamsSchema.parse({ tenantId, providerId, modelId });
         const provider = await this.getProviderById(validated.tenantId, validated.providerId, options);
         if (!provider) {
@@ -600,36 +661,74 @@ export class LLMService extends BaseAwsService {
         if (!model.isAvailable) {
             throw new ModelNotAvailableError(validated.providerId, validated.modelId);
         }
+        // Check budget and rate limits before proceeding
+        const usageTracker = getUsageTrackingService();
+        const userCredentials = options === null || options === void 0 ? void 0 : options.userCredentials;
+        const trackingParams = {
+            tenantId: validated.tenantId,
+            providerId: validated.providerId,
+            modelId: validated.modelId,
+            userCredentials,
+        };
+        await usageTracker.checkLimitsOrThrow(trackingParams);
+        // Strip reasoning content from input messages
+        // Reasoning is model output (internal thought process) and should never be sent back as input
+        // Models don't need their previous reasoning, and non-thinking models will reject it
+        let messagesToProcess = stripReasoningFromMessages(messages);
         // Bedrock has strict filename requirements for document attachments
         // Sanitize filenames before converting to model messages
-        const messagesToConvert = provider.type === "bedrock" ? sanitizeMessagesForBedrock(messages) : messages;
-        const modelMessages = convertToModelMessages(messagesToConvert);
+        if (provider.type === "bedrock") {
+            messagesToProcess = sanitizeMessagesForBedrock(messagesToProcess);
+        }
+        // Ensure tool inputs are never null/undefined when tools are used
+        // This prevents serialization issues across all providers, not just Bedrock
+        if (options === null || options === void 0 ? void 0 : options.tools) {
+            messagesToProcess = sanitizeToolPartsForBedrock(messagesToProcess);
+        }
+        // Normalize custom tool parts to standard toolInvocations
+        // This ensures convertToModelMessages can correctly process tool calls from history
+        messagesToProcess = normalizeToolParts(messagesToProcess);
+        // Convert messages to model format
+        // When tools are provided, pass them for proper tool call/result handling
+        const modelMessages = await convertToModelMessages(messagesToProcess, (options === null || options === void 0 ? void 0 : options.tools) ? { tools: options.tools } : undefined);
         // Get the AI SDK model instance using provider type and model identifier
         // Pass user credentials for Bedrock to enable session token authentication
-        const aiModel = this.getAISDKProvider(provider, options === null || options === void 0 ? void 0 : options.userCredentials).languageModel(model.id);
-        // Wrap with reasoning middleware if model supports thinking
-        const wrappedModel = model.capabilities.includes("thinking")
-            ? wrapLanguageModel({
-                model: aiModel,
-                middleware: extractReasoningMiddleware({
-                    tagName: "think",
-                }),
-            })
-            : aiModel;
+        const aiModel = this.getAISDKProvider(provider, options === null || options === void 0 ? void 0 : options.userCredentials).languageModel(model.modelId);
         this.logger.info("Streaming response", {
             providerId: provider.id,
             providerName: provider.name,
-            modelId: model.id,
+            modelId: model.modelId,
             modelName: model.name,
             messageCount: messages.length,
         });
-        return streamText({
-            model: wrappedModel,
-            system: systemPrompt,
-            messages: modelMessages,
-            maxOutputTokens: model.maxOutputTokens,
-            abortSignal: options === null || options === void 0 ? void 0 : options.abortSignal,
-        });
+        let budgetTokens = 0;
+        if (provider.type === "bedrock" && model.capabilities.includes("thinking")) {
+            budgetTokens = (_a = model.thinkingBudgetTokens) !== null && _a !== void 0 ? _a : THINKING_BUDGET.DEFAULT_TOKENS;
+            if (model.maxOutputTokens <= budgetTokens) {
+                throw new ValidationError(`Model maxOutputTokens (${model.maxOutputTokens}) must be greater than thinkingBudgetTokens (${budgetTokens}) for thinking models`);
+            }
+        }
+        const maxOutputTokens = model.maxOutputTokens - budgetTokens;
+        return streamText(Object.assign(Object.assign({ model: aiModel, system: systemPrompt, messages: modelMessages, maxOutputTokens, abortSignal: options === null || options === void 0 ? void 0 : options.abortSignal }, ((options === null || options === void 0 ? void 0 : options.tools) && {
+            tools: options.tools,
+            toolChoice: (_b = options.toolChoice) !== null && _b !== void 0 ? _b : "auto",
+        })), { providerOptions: Object.assign({}, (provider.type === "bedrock" && model.capabilities.includes("thinking")
+                ? {
+                    bedrock: {
+                        reasoningConfig: {
+                            type: "enabled",
+                            budgetTokens,
+                        },
+                    },
+                }
+                : undefined)), onFinish: async ({ usage, text }) => {
+                // Record usage for budget/rate limit tracking
+                await usageTracker.recordUsage(trackingParams, model, usage);
+                // Call onComplete callback if provided (for server-side persistence)
+                if (options === null || options === void 0 ? void 0 : options.onComplete) {
+                    await options.onComplete(text);
+                }
+            } }));
     }
     /**
      * Generate a complete text response from an AI model (non-streaming).
@@ -675,22 +774,47 @@ export class LLMService extends BaseAwsService {
         if (!model.isAvailable) {
             throw new ModelNotAvailableError(validated.providerId, validated.modelId);
         }
+        // Check budget and rate limits before proceeding
+        const usageTracker = getUsageTrackingService();
+        const userCredentials = options === null || options === void 0 ? void 0 : options.userCredentials;
+        const trackingParams = {
+            tenantId: validated.tenantId,
+            providerId: validated.providerId,
+            modelId: validated.modelId,
+            userCredentials,
+        };
+        await usageTracker.checkLimitsOrThrow(trackingParams);
         // Get the AI SDK model instance
         // Pass user credentials for Bedrock to enable session token authentication
-        const aiModel = this.getAISDKProvider(provider, options === null || options === void 0 ? void 0 : options.userCredentials).languageModel(model.id);
+        const aiModel = this.getAISDKProvider(provider, userCredentials).languageModel(model.modelId);
         this.logger.info("Generating text", {
             providerId: provider.id,
             providerName: provider.name,
-            modelId: model.id,
+            modelId: model.modelId,
             modelName: model.name,
             messageCount: messages.length,
         });
-        return generateText({
-            model: aiModel,
-            system: systemPrompt,
-            messages,
-            maxOutputTokens: model.maxOutputTokens,
-        });
+        try {
+            const result = await generateText({
+                model: aiModel,
+                system: systemPrompt,
+                messages,
+                maxOutputTokens: model.maxOutputTokens,
+            });
+            // Record usage for budget/rate limit tracking
+            await usageTracker.recordUsage(trackingParams, model, result.usage);
+            return result;
+        }
+        catch (error) {
+            // Try to extract and record usage from error if available
+            // Some AI SDK errors include partial usage data
+            if (error && typeof error === "object" && "usage" in error) {
+                if (error.usage) {
+                    await usageTracker.recordUsage(trackingParams, model, error.usage);
+                }
+            }
+            throw error;
+        }
     }
     // ============================================================================
     // Basic Assistant Default Model Methods
@@ -784,9 +908,13 @@ export class LLMService extends BaseAwsService {
             event: "basic_assistant_default_changed",
             tenantId,
             previousDefault: currentDefault
-                ? { providerId: currentDefault.model.providerId, modelId: currentDefault.model.id }
+                ? {
+                    providerId: currentDefault.model.providerId,
+                    id: currentDefault.model.id,
+                    modelId: currentDefault.model.modelId,
+                }
                 : null,
-            newDefault: { providerId, modelId },
+            newDefault: { providerId, id: modelId },
             timestamp: Date.now(),
         }));
         return updatedModel;
@@ -812,17 +940,19 @@ export class LLMService extends BaseAwsService {
         const defaultModels = models.filter((m) => m.isBasicAssistantDefault);
         // Clear all defaults (should only be one, but handle edge cases)
         for (const model of defaultModels) {
-            const builder = new UpdateExpressionBuilder()
+            const builder = this.db
+                .createUpdateBuilder()
                 .set("updatedAt", Date.now())
                 .set("isBasicAssistantDefault", false);
             const { UpdateExpression, ExpressionAttributeNames, ExpressionAttributeValues } = builder.build();
             await this.db.updateItem({
-                PK: buildTenantPK(tenantId),
-                SK: buildModelSK(model.providerId, model.id),
+                PK: llmProviderKeys.pk(tenantId),
+                SK: llmProviderKeys.modelSk(model.providerId, model.id),
             }, UpdateExpression, ExpressionAttributeNames, ExpressionAttributeValues, { userCredentials: options === null || options === void 0 ? void 0 : options.userCredentials });
             this.logger.debug("Cleared Basic Assistant default from model", {
                 tenantId,
-                modelId: model.id,
+                id: model.id,
+                modelId: model.modelId,
                 providerId: model.providerId,
             });
         }
@@ -871,7 +1001,8 @@ export class LLMService extends BaseAwsService {
             if (!model.isAvailable) {
                 this.logger.warn("Basic Assistant default model is unavailable", {
                     tenantId,
-                    modelId: model.id,
+                    id: model.id,
+                    modelId: model.modelId,
                     modelName: model.name,
                 });
                 return this.getFallbackBasicAssistantModel(tenantId, LLM_ERROR_MESSAGES.BASIC_ASSISTANT_MODEL_UNAVAILABLE, options);
@@ -907,13 +1038,15 @@ export class LLMService extends BaseAwsService {
                     tenantId,
                     providerId: provider.id,
                     providerName: provider.name,
-                    modelId: model.id,
+                    id: model.id,
+                    modelId: model.modelId,
                     modelName: model.name,
                     fallbackReason,
                 });
                 return {
                     provider: {
                         id: provider.id,
+                        createdById: provider.createdById,
                         name: provider.name,
                         type: provider.type,
                         description: provider.description,
@@ -956,16 +1089,7 @@ export function getLLMService() {
         llmServiceInstance = new LLMService({
             logger: LoggerFactory.createLogger({
                 serviceName: "llm",
-                userContext: {
-                    userId: "system",
-                    email: "system@captify.com",
-                    name: "System Service Account",
-                    roles: ["service"],
-                    metadata: {
-                        isSystemUser: true,
-                        note: "Service-level logger - user context should be provided by caller",
-                    },
-                },
+                userContext: USER_CONTEXT_WITH_METADATA,
             }),
         });
     }