feat(middleware): add custom middleware development guide - #119
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the WalkthroughIntroduces a full middleware system (types, registry, factory, built-in analytics), integrates it into BaseProvider to wrap AI SDK models, expands public exports, adds docs (middleware guides and Guardrails integration rewrite), and adds unit tests validating analytics middleware behavior. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant BaseProvider
participant MiddlewareFactory
participant Registry as MiddlewareRegistry
participant Chain as Middleware Chain
participant Model as AI SDK Model
participant Provider as LLM Provider
Client->>BaseProvider: generate(params, options)
BaseProvider->>MiddlewareFactory: createContext(provider, model, options, session)
BaseProvider->>MiddlewareFactory: applyMiddleware(Model, context, factoryOptions)
MiddlewareFactory->>Registry: buildChain(context, config)
Registry-->>MiddlewareFactory: LanguageModelV1Middleware[]
MiddlewareFactory-->>BaseProvider: Wrapped Model
BaseProvider->>Model: generate(...)
activate Chain
Chain->>Model: wrapGenerate(doGenerate)
Model->>Provider: request
Provider-->>Model: response (+usage)
Chain-->>BaseProvider: result (may include analytics)
deactivate Chain
BaseProvider-->>Client: result
sequenceDiagram
autonumber
participant Client
participant BaseProvider
participant MiddlewareFactory
participant Chain as Middleware Chain
participant Model as AI SDK Model
participant Provider as LLM Provider
Client->>BaseProvider: stream(params, options)
BaseProvider->>MiddlewareFactory: applyMiddleware(Model, context, factoryOptions)
MiddlewareFactory-->>BaseProvider: Wrapped Model
BaseProvider->>Model: stream(...)
activate Chain
Chain->>Model: wrapStream(doStream)
Model->>Provider: start stream
Provider-->>Model: streamed chunks
Chain-->>BaseProvider: streamed result (metadata via rawResponse)
deactivate Chain
BaseProvider-->>Client: stream handle
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
✨ Finishing Touches🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
5be6e2d to
2372e21
Compare
amreetkhuntia
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/lib/middleware/builtin/analytics.ts (1)
138-156: getAnalyticsMetrics/clearAnalyticsMetrics are stubs and can’t access per-request MaprequestMetrics is scoped inside createAnalyticsMiddleware(), so the exported getters always return an empty Map. Either remove these exports until a global registry exists or implement a module-level store.
-export function getAnalyticsMetrics(): Map<string, Record<string, unknown>> { - // This would need to be implemented with a global registry - // For now, return empty map - return new Map(); -} +// Module-level store shared by all middleware instances +const globalRequestMetrics = new Map<string, Record<string, unknown>>(); + +export function getAnalyticsMetrics(): Map<string, Record<string, unknown>> { + return globalRequestMetrics; +} -export function clearAnalyticsMetrics(): void { - // This would need to be implemented with a global registry - // For now, do nothing -} +export function clearAnalyticsMetrics(): void { + globalRequestMetrics.clear(); +}Additionally, update createAnalyticsMiddleware to write into globalRequestMetrics instead of a per-instance Map.
🧹 Nitpick comments (29)
docs/CUSTOM-MIDDLEWARE-GUIDE.md (7)
105-119: Add missing import for streamText in the usage example.Without importing
streamText, the example won’t run as-is.Apply this diff to the example:
import { wrapLanguageModel } from "ai"; +import { streamText } from "ai"; import { yourCustomMiddleware } from "./middleware/your-custom-middleware";
132-151: Avoid logging full prompts/results; redact or summarize to prevent PII leakage.The logging example prints prompts and timings verbatim, which risks exposing sensitive data in logs. Prefer summaries, hashes, or redaction.
[security]Suggested tweak:
- console.log( - `[${new Date().toISOString()}] Generate called with prompt:`, - typeof params.prompt === "string" - ? params.prompt.substring(0, 100) + "..." - : "Complex prompt", - ); + const promptPreview = + typeof params.prompt === "string" + ? `${params.prompt.length} chars (preview: ${params.prompt.slice(0, 40)}…)` + : "non-string prompt"; + console.log(`[${new Date().toISOString()}] Generate called`, { + prompt: promptPreview, + });
203-226: Use model-reported token usage when available instead of whitespace counts.You’re counting tokens via whitespace, which is inaccurate across tokenizers. The analytics snippet can mirror the built-in middleware by reading
result.usage.- // Calculate input tokens (simplified example) - const inputTokens = - typeof params.prompt === "string" ? params.prompt.split(/\s+/).length : 0; - - // Calculate output tokens (simplified example) - const outputTokens = result.text?.split(/\s+/).length || 0; + // Prefer model-reported usage metrics (fallback to 0 if unavailable) + const inputTokens = result.usage?.promptTokens ?? 0; + const outputTokens = result.usage?.completionTokens ?? 0;
85-97: Node/Web Streams typing note for TransformStream.In TS projects that don’t include the DOM lib,
TransformStreamwon’t be typed. Add an import fromstream/web(Node >=18) or ensure"lib": ["ES2022","DOM"]in tsconfig.Add to the top of the code block if needed:
import { TransformStream } from "stream/web";
167-177: Stabilize cache keys to avoid false misses.
JSON.stringifyis order-sensitive for objects and can yield different keys for semantically equal inputs. Consider a stable stringify or hashing strategy.
[performance]Example:
import stableStringify from "json-stable-stringify"; // … const cacheKey = stableStringify({ prompt: params.prompt, temperature: params.settings?.temperature, maxOutputTokens: params.settings?.maxOutputTokens, });
37-40: Minor list formatting.If this renders as “1… 2… 3…” on the same line in some markdown viewers, insert a blank line before the list or ensure each item ends with two spaces. Otherwise fine.
281-288: Tighten phrasing on chain flow.Consider adding a period and clarifying “reverse order back to the caller”.
-5. The response flows back through the chain in reverse order +5. The response flows back through the chain in reverse order to the caller.docs/GUARDRAILS-AI-INTEGRATION.md (1)
43-48: Import streamText in the usage example.The sample calls
streamTextwithout importing it.-import { wrapLanguageModel } from "ai"; +import { wrapLanguageModel, streamText } from "ai";src/lib/middleware/types.ts (4)
67-75: Prefer JsonValue for serializable context/options.
options: Record<string, unknown>prevents easy serialization in registries/telemetry. Since you already importJsonValue, using it improves consistency and downstream logging.- /** Request options */ - options: Record<string, unknown>; + /** Request options (serializable) */ + options: Record<string, JsonValue>;
53-56: Align MiddlewareConditions.options to JsonValue as well.Same rationale as context; it simplifies predicate checks and potential persistence.
- /** Apply only when certain options are present */ - options?: Record<string, unknown>; + /** Apply only when certain options are present (serializable) */ + options?: Record<string, JsonValue>;
120-127: Naming consistency: rateLimit vs rate limiting.Docs mention
createRateLimitingMiddleware, while the union here uses"rateLimit". Consider documenting the mapping or adding"rateLimiting"as an alias to reduce friction.
144-162: Optional: consider readonly props in MiddlewareFactoryOptions.global.If these are never mutated after construction, marking them as
readonlyconveys intent. Not required; just a type-safety polish.src/lib/index.ts (1)
55-63: Consider exporting additional middleware types for custom dev ergonomics.Expose metadata and config-related types so users can type their custom middleware and presets without deep imports.
export type { NeuroLinkMiddleware, MiddlewareContext, MiddlewareFactoryOptions, -} from "./middleware/types.js"; + // Optional: additional helpful types + NeuroLinkMiddlewareMetadata, + MiddlewareConfig, + MiddlewareConditions, + MiddlewareRegistrationOptions, + MiddlewareExecutionResult, + MiddlewareChainStats, + BuiltInMiddlewareType, +} from "./middleware/types.js";test/middleware-basic.test.ts (1)
4-4: Remove unused type import.
MiddlewareContextisn’t used in this file.-import type { MiddlewareContext } from "../src/lib/middleware/types.js";src/lib/core/baseProvider.ts (3)
628-642: Skipping middleware when tools are disabled is too coarseCurrently, shouldSkipMiddleware() returns true when disableTools is true. Many middleware (analytics, logging, guardrails) are orthogonal to tool execution and should still apply.
- Suggestion: Only skip middleware explicitly marked as tool-dependent (e.g., via metadata or config), or introduce a per-middleware condition to respect disableTools. Keep disableMiddleware as the global off switch.
587-623: Middleware options exist only as “loose” fields on options; add types for safetyextractMiddlewareOptions pulls middlewareConfig, enabledMiddleware, disabledMiddleware, middlewarePreset from the options bag using Record<string, unknown>. These fields aren’t defined on TextGenerationOptions or StreamOptions, so callers don’t get type help and can mistype keys.
- Recommendation: Extend TextGenerationOptions and StreamOptions with an optional middleware field (or the individual fields) to provide compile-time safety and autocomplete.
- Alternative: Export a typed helper to build MiddlewareFactoryOptions and document the supported keys.
716-724: Be cautious logging tool params/results — potential PII and large payloadsThe debug logs include params and possibly large result objects. This can leak sensitive data and bloat logs.
- Suggestion: Truncate large strings/objects and redact obvious PII keys (e.g., apiKey, password, token). Gate verbose payload logging behind a dedicated flag.
- logger.debug(`[BaseProvider] Tool execution successful: ${toolName}`, { - resultType: typeof result, - hasResult: result !== null && result !== undefined, - toolName, - }); + logger.debug(`[BaseProvider] Tool execution successful: ${toolName}`, { + resultType: typeof result, + hasResult: result !== null && result !== undefined, + sample: typeof result === "string" ? result.slice(0, 200) : undefined, + toolName, + });Also applies to: 730-735, 743-749
docs/MIDDLEWARE.md (2)
476-482: Tighten preset list formatting for readabilityMinor doc polish: use a standard bulleted list without backticks around preset names to avoid grammar warnings and improve readability.
-- `development`: Logging and basic analytics for development -- `production`: Optimized for production with caching and rate limiting -- `security`: Enhanced security with guardrails and monitoring -- `performance`: Optimized for performance with caching and retries -- `enterprise`: Full enterprise feature set with all middleware -- `minimal`: Minimal overhead with only essential features +• development — Logging and basic analytics for development +• production — Optimized for production with caching and rate limiting +• security — Enhanced security with guardrails and monitoring +• performance — Optimized for performance with caching and retries +• enterprise — Full enterprise feature set with all middleware +• minimal — Minimal overhead with only essential features
404-416: Sample import paths: align naming with repo structureExamples reference “custom.js” and “middleware” paths that may not exist in this repo. Consider aligning examples with actual paths (e.g., builtin/analytics) or add a short note that paths are illustrative.
src/lib/middleware/builtin/analytics.ts (2)
24-80: Good middleware shape and logging; consider including model/provider metadataSolid timing/usage tracking with clear logs. To make analytics more actionable, include provider/model identifiers if available via params or model to help downstream aggregation.
82-129: Inconsistent attachment points for analytics between generate and streamwrapGenerate attaches analytics at experimental_providerMetadata.neurolink.analytics, but wrapStream uses rawResponse.neurolink.analytics. This divergence complicates consumers.
- Recommendation: Standardize attachment location (prefer experimental_providerMetadata.neurolink.analytics for both). If stream shape prevents that, document the difference and add a helper to normalize extraction.
- if (!updatedResult.rawResponse) { updatedResult.rawResponse = {}; } - if (!updatedResult.rawResponse.neurolink) { - updatedResult.rawResponse.neurolink = {}; - } - updatedResult.rawResponse.neurolink.analytics = streamAnalytics; + if (!updatedResult.experimental_providerMetadata) { + updatedResult.experimental_providerMetadata = {}; + } + if (!updatedResult.experimental_providerMetadata.neurolink) { + updatedResult.experimental_providerMetadata.neurolink = {}; + } + updatedResult.experimental_providerMetadata.neurolink.analytics = streamAnalytics;src/lib/middleware/index.ts (1)
41-47: Re-export built-in analytics for DXGiven src/lib/middleware/builtin/analytics.ts exists, consider re-exporting createAnalyticsMiddleware here so consumers don’t need to know the internal folder layout.
-// export { analyticsMiddleware } from './built-in/analytics.js'; +export { createAnalyticsMiddleware } from "./builtin/analytics.js";src/lib/middleware/factory.ts (7)
93-99: Log when an unknown preset is requested.Currently, an invalid preset is silently ignored. Warn to aid diagnostics.
- if (options.preset) { - const presetConfig = this.getPresetConfig(options.preset); - if (presetConfig) { - Object.assign(config, presetConfig); - } - } + if (options.preset) { + const presetConfig = this.getPresetConfig(options.preset); + if (presetConfig) { + Object.assign(config, presetConfig); + } else { + logger.warn("Unknown middleware preset", { preset: options.preset }); + } + }
213-215: Use crypto.randomUUID() and avoid deprecated substr for requestId.Math.random and substr are suboptimal. Prefer a UUID and reuse one timestamp for consistency.
- timestamp: Date.now(), - requestId: `${provider}-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`, + timestamp: Date.now(), + requestId: `${provider}-${crypto.randomUUID()}`,Add this import at the top of the file:
+import { crypto } from "node:crypto";If you’d rather not import, on Node 20+ globalThis.crypto.randomUUID() also works.
338-353: Applied count should reflect the built chain, not availability of historical stats.appliedMiddleware is derived from aggregated stats presence, which can undercount newly-enabled middleware with no prior executions. Use chain.length and keep results based on stats.
- let totalExecutionTime = 0; - let appliedMiddleware = 0; + let totalExecutionTime = 0; + const appliedMiddleware = chain.length; @@ - if (config[middlewareId]?.enabled) { + if (config[middlewareId]?.enabled) { results[middlewareId] = { applied: true, executionTime: middlewareStats.averageExecutionTime, }; totalExecutionTime += middlewareStats.averageExecutionTime; - appliedMiddleware++; }Please confirm intended semantics: “applied” = included in the built chain for this context. If you instead want “executed historically,” consider renaming the field to avoid confusion.
371-379: Merge enabled/disabled lists across defaults and per-call overrides.You already deep-merge middlewareConfig. Do the same for enabledMiddleware/disabledMiddleware to make overrides additive.
- const _mergedOptions = { + const _mergedOptions: MiddlewareFactoryOptions = { ...defaultOptions, ...options, + enabledMiddleware: [ + ...(defaultOptions.enabledMiddleware ?? []), + ...(options.enabledMiddleware ?? []), + ], + disabledMiddleware: [ + ...(defaultOptions.disabledMiddleware ?? []), + ...(options.disabledMiddleware ?? []), + ], middlewareConfig: { ...defaultOptions.middlewareConfig, ...options.middlewareConfig, }, };
282-321: Derive preset middleware lists from getBuiltInPresets to avoid drift.The middleware arrays here duplicate getBuiltInPresets and can desync. Programmatically read keys from the presets map.
static getAvailablePresets(): Array<{ name: string; description: string; middleware: string[]; }> { - return [ + const presets = this.getBuiltInPresets(); + return [ { name: "development", description: "Logging and basic analytics for development", - middleware: ["logging", "analytics"], + middleware: Object.keys(presets.development), }, { name: "production", description: "Optimized for production with caching and rate limiting", - middleware: ["analytics", "caching", "rateLimit", "retry"], + middleware: Object.keys(presets.production), }, { name: "security", description: "Enhanced security with guardrails and monitoring", - middleware: ["guardrails", "logging", "rateLimit"], + middleware: Object.keys(presets.security), }, { name: "performance", description: "Optimized for performance with caching and retries", - middleware: ["caching", "retry", "timeout"], + middleware: Object.keys(presets.performance), }, { name: "enterprise", description: "Full enterprise feature set with all middleware", - middleware: [ - "analytics", - "guardrails", - "logging", - "caching", - "rateLimit", - "retry", - "timeout", - ], + middleware: Object.keys(presets.enterprise), }, { name: "minimal", description: "Minimal overhead with only essential features", - middleware: ["analytics"], + middleware: Object.keys(presets.minimal), }, ];
74-132: Add tests for config building and edge cases.Recommend unit tests for:
- Unknown middleware ID in options.middlewareConfig does not throw and deep-merges nested config.
- enabledMiddleware/disabledMiddleware merging (if you adopt the additive merge).
- Preset not found emits a warning (if you add the warn).
I can scaffold these with Vitest/Jest if you want test stubs wired to middlewareRegistry.buildChain mocks.
360-371: Ensure backward-compatible context support in createModelFactoryWe didn’t find any internal calls to
createModelFactoryin the repo, so changing its signature would only break external consumers. To add thecontextparameter without forcing all downstream callers to update, I recommend introducing a new overload/helper and deprecating the old one:• File: src/lib/middleware/factory.ts, around lines 360–371
• Keep the existingstatic createModelFactory(baseModelFactory: () => Promise<LanguageModelV1>, …)unchanged
• Add a new method:/** * @deprecated Use createContextualModelFactory for full context support. */ static createModelFactory( baseModelFactory: (context: MiddlewareContext) => Promise<LanguageModelV1>, defaultOptions: MiddlewareFactoryOptions = {}, ) { return async ( context: MiddlewareContext, options: MiddlewareFactoryOptions = {}, ): Promise<LanguageModelV1> => { const baseModel = await baseModelFactory(context); // Merge options… }; } /** * Creates a middleware-enabled model factory that receives context. */ static createContextualModelFactory( baseModelFactory: (context: MiddlewareContext) => Promise<LanguageModelV1>, defaultOptions: MiddlewareFactoryOptions = {}, ) { return async ( context: MiddlewareContext, options: MiddlewareFactoryOptions = {}, ): Promise<LanguageModelV1> => { const baseModel = await baseModelFactory(context); // Merge options… }; }This approach preserves the existing API, avoids breaking current users, and offers a clear migration path.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (11)
docs/CUSTOM-MIDDLEWARE-GUIDE.md(1 hunks)docs/GUARDRAILS-AI-INTEGRATION.md(1 hunks)docs/MIDDLEWARE.md(1 hunks)src/lib/core/baseProvider.ts(4 hunks)src/lib/index.ts(1 hunks)src/lib/middleware/builtin/analytics.ts(1 hunks)src/lib/middleware/factory.ts(1 hunks)src/lib/middleware/index.ts(1 hunks)src/lib/middleware/registry.ts(1 hunks)src/lib/middleware/types.ts(1 hunks)test/middleware-basic.test.ts(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (7)
src/lib/middleware/factory.ts (3)
src/lib/middleware/types.ts (4)
MiddlewareContext(61-75)MiddlewareConfig(35-42)MiddlewareChainStats(106-115)MiddlewareExecutionResult(92-101)src/lib/middleware/registry.ts (1)
middlewareRegistry(406-406)src/lib/utils/logger.ts (2)
logger(198-237)error(134-136)
src/lib/middleware/types.ts (2)
src/lib/index.ts (3)
NeuroLinkMiddleware(57-57)MiddlewareContext(58-58)MiddlewareFactoryOptions(59-59)src/lib/middleware/index.ts (11)
NeuroLinkMiddleware(20-20)LanguageModelV1Middleware(33-33)MiddlewareConfig(21-21)MiddlewareConditions(23-23)MiddlewareContext(22-22)MiddlewareRegistrationOptions(24-24)MiddlewareExecutionResult(25-25)MiddlewareChainStats(26-26)BuiltInMiddlewareType(29-29)MiddlewarePreset(27-27)MiddlewareFactoryOptions(28-28)
src/lib/middleware/builtin/analytics.ts (2)
src/lib/middleware/types.ts (2)
NeuroLinkMiddleware(27-30)NeuroLinkMiddlewareMetadata(8-21)src/lib/utils/logger.ts (2)
logger(198-237)error(134-136)
src/lib/middleware/registry.ts (2)
src/lib/middleware/types.ts (5)
NeuroLinkMiddleware(27-30)MiddlewareExecutionResult(92-101)MiddlewareRegistrationOptions(80-87)MiddlewareContext(61-75)MiddlewareConfig(35-42)src/lib/utils/logger.ts (2)
logger(198-237)error(134-136)
src/lib/middleware/index.ts (3)
src/lib/middleware/types.ts (2)
NeuroLinkMiddleware(27-30)MiddlewareConfig(35-42)src/lib/middleware/registry.ts (1)
middlewareRegistry(406-406)src/lib/middleware/factory.ts (2)
getAvailablePresets(277-322)MiddlewareFactory(16-385)
test/middleware-basic.test.ts (2)
src/lib/middleware/builtin/analytics.ts (1)
createAnalyticsMiddleware(12-136)src/lib/middleware/factory.ts (1)
MiddlewareFactory(16-385)
src/lib/core/baseProvider.ts (5)
src/lib/core/types.ts (1)
TextGenerationOptions(165-190)src/lib/types/streamTypes.ts (1)
StreamOptions(80-138)src/lib/middleware/factory.ts (1)
MiddlewareFactory(16-385)src/lib/utils/logger.ts (2)
logger(198-237)error(134-136)src/lib/middleware/types.ts (2)
MiddlewareFactoryOptions(144-162)MiddlewareConfig(35-42)
🪛 LanguageTool
docs/CUSTOM-MIDDLEWARE-GUIDE.md
[grammar] ~37-~37: There might be a mistake here.
Context: ...re they are passed to the language model 2. wrapGenerate: Wraps the doGenerate method of the l...
(QB_NEW_EN)
[grammar] ~38-~38: There might be a mistake here.
Context: ...doGeneratemethod of the language model 3.wrapStream: Wraps the doStream` method of the lan...
(QB_NEW_EN)
[grammar] ~286-~286: There might be a mistake here.
Context: ...ilteringMiddleware` filters the response 5. The response flows back through the chai...
(QB_NEW_EN)
docs/MIDDLEWARE.md
[style] ~34-~34: Three successive sentences begin with the same word. Consider rewording the sentence or use a thesaurus to find a synonym.
Context: ...D --> C C --> B B --> A ``` ## Middleware Types and Interfaces ### LanguageModel...
(ENGLISH_WORD_REPEAT_BEGINNING_RULE)
[grammar] ~476-~476: There might be a mistake here.
Context: ...ging and basic analytics for development - production: Optimized for production with caching ...
(QB_NEW_EN)
[grammar] ~477-~477: There might be a mistake here.
Context: ...roduction with caching and rate limiting - security: Enhanced security with guardrails and ...
(QB_NEW_EN)
[grammar] ~478-~478: There might be a mistake here.
Context: ... security with guardrails and monitoring - performance: Optimized for performance with caching...
(QB_NEW_EN)
[grammar] ~479-~479: There might be a mistake here.
Context: ...for performance with caching and retries - enterprise: Full enterprise feature set with all m...
(QB_NEW_EN)
[grammar] ~480-~480: There might be a mistake here.
Context: ...terprise feature set with all middleware - minimal: Minimal overhead with only essential f...
(QB_NEW_EN)
🔇 Additional comments (8)
src/lib/middleware/types.ts (1)
1-31: Solid foundational typing for middleware + metadata.Interfaces map cleanly to usage in analytics/factory and look extensible. The
readonly metadataonNeuroLinkMiddlewareis a good guardrail.src/lib/index.ts (1)
55-63: Good addition: cleanly exposes middleware types and factories.Public API surface is well-scoped. Exporting the factory and analytics helper will enable quick starts.
test/middleware-basic.test.ts (1)
6-10: AI SDK mock is fine; ensure CI has the alias.Mocking
"ai"avoids resolution issues where the SDK isn’t installed. Keep this, or alternatively add a dev dependency onaifor type-only usage.docs/MIDDLEWARE.md (1)
140-199: Analytics example: clarify where analytics is attached on resultsDocs show attaching analytics in wrapGenerate to result.analytics. Elsewhere, the code attaches to experimental_providerMetadata.neurolink.analytics. Align the doc snippet to match the current implementation to avoid confusion for users reading result shapes.
src/lib/middleware/registry.ts (1)
146-163: Enabled logic: treat “defaultEnabled: true” consistently when config entry is absentCurrently, if a middleware has defaultEnabled: true and there’s no config entry, it’s applied (good). Ensure this behavior is documented so callers know they must explicitly disable such middleware via disabledMiddleware or config.enabled=false.
src/lib/middleware/index.ts (1)
95-101: Good: centralized access to aggregated statsThe convenience wrapper over registry stats is a nice touch and keeps the API cohesive.
src/lib/middleware/factory.ts (2)
20-72: Middleware wrapping flow is sound; graceful fallback on errors.The overall control flow in applyMiddleware is clear: build config → build chain → early-return on empty → wrap → time/trace → catch-and-fallback to base model. Good defensive posture and observability.
48-53: Action: resolved — keep using wrapLanguageModelVerified results:
- package.json declares ai@4.3.16.
- src/lib/middleware/factory.ts imports and uses wrapLanguageModel (import at top; usage around lines 48–52).
- node_modules/ai exports wrapLanguageModel; experimental_wrapLanguageModel exists only as a deprecated alias that points to wrapLanguageModel (typings and compiled files show both, with experimental marked deprecated).
Conclusion: the current import/usage of wrapLanguageModel is correct for the installed SDK version — no change required.
| import type { LanguageModelV2Middleware } from "@ai-sdk/provider"; | ||
|
|
||
| ### Output Guard (Post-processing) | ||
| export const yourGuardrailMiddleware: LanguageModelV2Middleware = { | ||
| wrapGenerate: async ({ doGenerate }) => { | ||
| const { text, ...rest } = await doGenerate(); | ||
|
|
||
| The Output Guard is responsible for validating the LLM's response before it is sent to the user. This can include checks for: | ||
| // Filtering approach, e.g., for PII or other sensitive information: | ||
| const cleanedText = text?.replace(/badword/g, "<REDACTED>"); | ||
|
|
||
| - **Toxicity Detection**: Ensure that the response is not toxic or offensive. | ||
| - **Fact-checking**: Verify that the information in the response is accurate. | ||
| - **Format Validation**: Ensure that the response is in the correct format (e.g., JSON, XML). | ||
| return { text: cleanedText, ...rest }; | ||
| }, | ||
|
|
||
| ## Structured Data Generation | ||
| // Note: Streaming guardrails are more complex to implement, | ||
| // as you do not have the full content of the stream until it's finished. | ||
| // A similar logic would be applied inside a transform stream for wrapStream. | ||
| }; | ||
| ``` |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Docs/API mismatch: using LanguageModelV2Middleware (provider) while the codebase standardizes on V1 from "ai".
The rest of the repo (types, built-ins, tests) uses LanguageModelV1Middleware from the ai package. Mixing V2 here will confuse implementers and won’t match the exported types.
Apply this diff:
-import type { LanguageModelV2Middleware } from "@ai-sdk/provider";
+import type { LanguageModelV1Middleware } from "ai";
-export const yourGuardrailMiddleware: LanguageModelV2Middleware = {
+export const yourGuardrailMiddleware: LanguageModelV1Middleware = {
wrapGenerate: async ({ doGenerate }) => {
const { text, ...rest } = await doGenerate();🏁 Script executed:
#!/bin/bash
# Verify no other references to V2 middleware remain
rg -nP "LanguageModelV2Middleware|@ai-sdk/provider" -C2Length of output: 7060
Align docs with V1 middleware standard
Verified that the codebase exclusively uses LanguageModelV1Middleware from the ai package and no other references to LanguageModelV2Middleware remain outside of the docs. Update the Guardrails integration guide accordingly.
• File: docs/GUARDRAILS-AI-INTEGRATION.md (lines 55–58)
Replace the V2 import and type with the V1 equivalents.
-import type { LanguageModelV2Middleware } from "@ai-sdk/provider";
+import type { LanguageModelV1Middleware } from "ai";
-export const yourGuardrailMiddleware: LanguageModelV2Middleware = {
+export const yourGuardrailMiddleware: LanguageModelV1Middleware = {
wrapGenerate: async ({ doGenerate }) => {
const { text, ...rest } = await doGenerate();
// … 📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| import type { LanguageModelV2Middleware } from "@ai-sdk/provider"; | |
| ### Output Guard (Post-processing) | |
| export const yourGuardrailMiddleware: LanguageModelV2Middleware = { | |
| wrapGenerate: async ({ doGenerate }) => { | |
| const { text, ...rest } = await doGenerate(); | |
| The Output Guard is responsible for validating the LLM's response before it is sent to the user. This can include checks for: | |
| // Filtering approach, e.g., for PII or other sensitive information: | |
| const cleanedText = text?.replace(/badword/g, "<REDACTED>"); | |
| - **Toxicity Detection**: Ensure that the response is not toxic or offensive. | |
| - **Fact-checking**: Verify that the information in the response is accurate. | |
| - **Format Validation**: Ensure that the response is in the correct format (e.g., JSON, XML). | |
| return { text: cleanedText, ...rest }; | |
| }, | |
| ## Structured Data Generation | |
| // Note: Streaming guardrails are more complex to implement, | |
| // as you do not have the full content of the stream until it's finished. | |
| // A similar logic would be applied inside a transform stream for wrapStream. | |
| }; | |
| ``` | |
| import type { LanguageModelV1Middleware } from "ai"; | |
| export const yourGuardrailMiddleware: LanguageModelV1Middleware = { | |
| wrapGenerate: async ({ doGenerate }) => { | |
| const { text, ...rest } = await doGenerate(); | |
| // Filtering approach, e.g., for PII or other sensitive information: | |
| const cleanedText = text?.replace(/badword/g, "<REDACTED>"); | |
| return { text: cleanedText, ...rest }; | |
| }, | |
| // Note: Streaming guardrails are more complex to implement, | |
| // as you do not have the full content of the stream until it's finished. | |
| // A similar logic would be applied inside a transform stream for wrapStream. | |
| }; |
🤖 Prompt for AI Agents
In docs/GUARDRAILS-AI-INTEGRATION.md around lines 55 to 71, the example imports
and types LanguageModelV2Middleware from "@ai-sdk/provider" but the codebase
uses LanguageModelV1Middleware from the "ai" package; update the import to pull
LanguageModelV1Middleware from "ai" and change the exported const type to
LanguageModelV1Middleware, ensuring any type-specific params or method names in
the example match the V1 middleware shape.
| /** | ||
| * Get AI SDK model with middleware applied | ||
| * This method wraps the base model with any configured middleware | ||
| */ | ||
| protected async getAISDKModelWithMiddleware( | ||
| options: TextGenerationOptions | StreamOptions = {}, | ||
| ): Promise<LanguageModelV1> { | ||
| // Get the base model | ||
| const baseModel = await this.getAISDKModel(); | ||
|
|
||
| // Check if middleware should be applied | ||
| const middlewareOptions = this.extractMiddlewareOptions(options); | ||
| if (!middlewareOptions || this.shouldSkipMiddleware(options)) { | ||
| return baseModel; | ||
| } | ||
|
|
||
| try { | ||
| // Create middleware context | ||
| const context = MiddlewareFactory.createContext( | ||
| this.providerName, | ||
| this.modelName, | ||
| options as Record<string, unknown>, | ||
| { | ||
| sessionId: this.sessionId, | ||
| userId: this.userId, | ||
| }, | ||
| ); | ||
|
|
||
| // Apply middleware to the model | ||
| const wrappedModel = MiddlewareFactory.applyMiddleware( | ||
| baseModel, | ||
| context, | ||
| middlewareOptions, | ||
| ); | ||
|
|
||
| logger.debug(`Applied middleware to ${this.providerName} model`, { | ||
| provider: this.providerName, | ||
| model: this.modelName, | ||
| hasMiddleware: true, | ||
| }); | ||
|
|
||
| return wrappedModel; | ||
| } catch (error) { | ||
| logger.warn( | ||
| `Failed to apply middleware to ${this.providerName}, using base model`, | ||
| { | ||
| error: error instanceof Error ? error.message : String(error), | ||
| }, | ||
| ); | ||
|
|
||
| // Return base model on middleware failure to maintain functionality | ||
| return baseModel; | ||
| } | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
getAISDKModelWithMiddleware is never invoked — middleware currently won’t run
You added a solid wrapper with fallback, but generate() still calls getAISDKModel(), so no middleware is ever applied in the core path.
- Action: In generate(), obtain the model via getAISDKModelWithMiddleware(options) so the chain can intercept generate/stream behavior.
- Note: Consider documenting in provider subclasses that executeStream implementations should also use getAISDKModelWithMiddleware when building the AI SDK model.
Example change outside this hunk (call site in generate):
// Before
const model = await this.getAISDKModel();
// After
const model = await this.getAISDKModelWithMiddleware(options);There was a problem hiding this comment.
Yes will addd this with optional thing. Knowingly didn't kept this
There was a problem hiding this comment.
@amreetkhuntia, understood! Building the infrastructure first before wiring it up is a solid approach. When you're ready to add the optional invocation, the integration point in generate() would be straightforward - just switching from getAISDKModel() to getAISDKModelWithMiddleware(options) based on your optional flag.
Let me know if you'd like any suggestions on the integration pattern when you get to that phase!
| // Apply explicit middleware configurations | ||
| if (options.middlewareConfig) { | ||
| for (const [middlewareId, middlewareConfig] of Object.entries( | ||
| options.middlewareConfig, | ||
| )) { | ||
| config[middlewareId] = { | ||
| ...config[middlewareId], | ||
| ...middlewareConfig, | ||
| }; | ||
| } | ||
| } |
There was a problem hiding this comment.
Fix runtime crash when overriding config for unknown middleware IDs.
If options.middlewareConfig contains an ID that wasn’t in the registry or a preset, config[middlewareId] is undefined and the spread operator will throw “Cannot convert undefined or null to object” at runtime.
Patch merges safely and deep-merges the nested config bag:
- if (options.middlewareConfig) {
- for (const [middlewareId, middlewareConfig] of Object.entries(
- options.middlewareConfig,
- )) {
- config[middlewareId] = {
- ...config[middlewareId],
- ...middlewareConfig,
- };
- }
- }
+ if (options.middlewareConfig) {
+ for (const [middlewareId, override] of Object.entries(
+ options.middlewareConfig,
+ )) {
+ const existing: MiddlewareConfig = config[middlewareId] ?? {
+ enabled: false,
+ config: {},
+ };
+ config[middlewareId] = {
+ ...existing,
+ ...override,
+ // ensure nested config is merged instead of clobbered
+ config: {
+ ...(existing.config ?? {}),
+ ...(override.config ?? {}),
+ },
+ };
+ }
+ }Also consider adding a unit test to assert this path with an unknown ID and nested config.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // Apply explicit middleware configurations | |
| if (options.middlewareConfig) { | |
| for (const [middlewareId, middlewareConfig] of Object.entries( | |
| options.middlewareConfig, | |
| )) { | |
| config[middlewareId] = { | |
| ...config[middlewareId], | |
| ...middlewareConfig, | |
| }; | |
| } | |
| } | |
| // Apply explicit middleware configurations | |
| if (options.middlewareConfig) { | |
| for (const [middlewareId, override] of Object.entries( | |
| options.middlewareConfig, | |
| )) { | |
| const existing: MiddlewareConfig = config[middlewareId] ?? { | |
| enabled: false, | |
| config: {}, | |
| }; | |
| config[middlewareId] = { | |
| ...existing, | |
| ...override, | |
| // ensure nested config is merged instead of clobbered | |
| config: { | |
| ...(existing.config ?? {}), | |
| ...(override.config ?? {}), | |
| }, | |
| }; | |
| } | |
| } |
🤖 Prompt for AI Agents
In src/lib/middleware/factory.ts around lines 101 to 111, the code blindly
spreads config[middlewareId] which can be undefined when
options.middlewareConfig contains an unknown middleware ID; update the logic to
ensure config[middlewareId] is initialized to an object before merging and
perform a deep/recursive merge of the nested config bag (e.g., use a safe deep
merge utility or implement a small recursive merge) so overrides don't throw
when baseline is missing; also add a unit test that passes an unknown middleware
ID with nested config to assert it initializes and deep-merges correctly.
| const mergedConfig = { | ||
| ...globalConfig, | ||
| ...config?.config, | ||
| }; | ||
|
|
||
| // Create wrapper that tracks execution | ||
| const wrappedMiddleware: LanguageModelV1Middleware = {}; |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Merged middleware config is computed but never used
mergedConfig is created and then discarded; no configuration actually reaches middleware. This makes enabled/config entries largely moot beyond on/off.
- Action: Thread mergedConfig into the middleware call path (e.g., via args.params.__neurolink.middlewareConfig[middlewareId]) so middleware can read it.
🤖 Prompt for AI Agents
In src/lib/middleware/registry.ts around lines 210 to 216, mergedConfig is
computed but never passed into middleware; update the middleware invocation path
to attach mergedConfig into
args.params.__neurolink.middlewareConfig[middlewareId] so the middleware can
read its config. Ensure you create or reuse args.params.__neurolink and its
middlewareConfig object safely (non-destructive to other params), assign
mergedConfig under the current middlewareId before calling the wrapped
middleware, and clean up or copy as needed to avoid mutating external objects
unexpectedly.
| if (middleware.transformParams) { | ||
| wrappedMiddleware.transformParams = async (args) => { | ||
| const startTime = Date.now(); | ||
| try { | ||
| const result = await middleware.transformParams!(args); | ||
| this.recordExecution(middleware.metadata.id, startTime, true); | ||
| return result; | ||
| } catch (error) { | ||
| this.recordExecution( | ||
| middleware.metadata.id, | ||
| startTime, | ||
| false, | ||
| error as Error, | ||
| ); | ||
| throw error; | ||
| } | ||
| }; | ||
| } | ||
|
|
||
| if (middleware.wrapGenerate) { | ||
| wrappedMiddleware.wrapGenerate = async (args) => { | ||
| const startTime = Date.now(); | ||
| try { | ||
| const result = await middleware.wrapGenerate!(args); | ||
| this.recordExecution(middleware.metadata.id, startTime, true); | ||
| return result; | ||
| } catch (error) { | ||
| this.recordExecution( | ||
| middleware.metadata.id, | ||
| startTime, | ||
| false, | ||
| error as Error, | ||
| ); | ||
| throw error; | ||
| } | ||
| }; | ||
| } | ||
|
|
||
| if (middleware.wrapStream) { | ||
| wrappedMiddleware.wrapStream = async (args) => { | ||
| const startTime = Date.now(); | ||
| try { | ||
| const result = await middleware.wrapStream!(args); | ||
| this.recordExecution(middleware.metadata.id, startTime, true); | ||
| return result; | ||
| } catch (error) { | ||
| this.recordExecution( | ||
| middleware.metadata.id, | ||
| startTime, | ||
| false, | ||
| error as Error, | ||
| ); | ||
| throw error; | ||
| } | ||
| }; | ||
| } | ||
|
|
||
| return wrappedMiddleware; | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Inject config into args and record stats around the wrapped calls
Pass the merged configuration to middleware by augmenting params; keep the existing metrics recording.
// Create wrapper that tracks execution
const wrappedMiddleware: LanguageModelV1Middleware = {};
if (middleware.transformParams) {
wrappedMiddleware.transformParams = async (args) => {
const startTime = Date.now();
try {
- const result = await middleware.transformParams!(args);
+ const nextArgs = {
+ ...args,
+ params: {
+ ...args.params,
+ __neurolink: {
+ ...(args.params?.__neurolink ?? {}),
+ middlewareConfig: {
+ ...(args.params?.__neurolink?.middlewareConfig ?? {}),
+ [middleware.metadata.id]: mergedConfig,
+ },
+ },
+ },
+ };
+ const result = await middleware.transformParams!(nextArgs);
this.recordExecution(middleware.metadata.id, startTime, true);
return result;
} catch (error) {
this.recordExecution(
middleware.metadata.id,
startTime,
false,
error as Error,
);
throw error;
}
};
}
if (middleware.wrapGenerate) {
wrappedMiddleware.wrapGenerate = async (args) => {
const startTime = Date.now();
try {
- const result = await middleware.wrapGenerate!(args);
+ const nextArgs = {
+ ...args,
+ params: {
+ ...args.params,
+ __neurolink: {
+ ...(args.params?.__neurolink ?? {}),
+ middlewareConfig: {
+ ...(args.params?.__neurolink?.middlewareConfig ?? {}),
+ [middleware.metadata.id]: mergedConfig,
+ },
+ },
+ },
+ };
+ const result = await middleware.wrapGenerate!(nextArgs);
this.recordExecution(middleware.metadata.id, startTime, true);
return result;
} catch (error) {
this.recordExecution(
middleware.metadata.id,
startTime,
false,
error as Error,
);
throw error;
}
};
}
if (middleware.wrapStream) {
wrappedMiddleware.wrapStream = async (args) => {
const startTime = Date.now();
try {
- const result = await middleware.wrapStream!(args);
+ const nextArgs = {
+ ...args,
+ params: {
+ ...args.params,
+ __neurolink: {
+ ...(args.params?.__neurolink ?? {}),
+ middlewareConfig: {
+ ...(args.params?.__neurolink?.middlewareConfig ?? {}),
+ [middleware.metadata.id]: mergedConfig,
+ },
+ },
+ },
+ };
+ const result = await middleware.wrapStream!(nextArgs);
this.recordExecution(middleware.metadata.id, startTime, true);
return result;
} catch (error) {
this.recordExecution(
middleware.metadata.id,
startTime,
false,
error as Error,
);
throw error;
}
};
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if (middleware.transformParams) { | |
| wrappedMiddleware.transformParams = async (args) => { | |
| const startTime = Date.now(); | |
| try { | |
| const result = await middleware.transformParams!(args); | |
| this.recordExecution(middleware.metadata.id, startTime, true); | |
| return result; | |
| } catch (error) { | |
| this.recordExecution( | |
| middleware.metadata.id, | |
| startTime, | |
| false, | |
| error as Error, | |
| ); | |
| throw error; | |
| } | |
| }; | |
| } | |
| if (middleware.wrapGenerate) { | |
| wrappedMiddleware.wrapGenerate = async (args) => { | |
| const startTime = Date.now(); | |
| try { | |
| const result = await middleware.wrapGenerate!(args); | |
| this.recordExecution(middleware.metadata.id, startTime, true); | |
| return result; | |
| } catch (error) { | |
| this.recordExecution( | |
| middleware.metadata.id, | |
| startTime, | |
| false, | |
| error as Error, | |
| ); | |
| throw error; | |
| } | |
| }; | |
| } | |
| if (middleware.wrapStream) { | |
| wrappedMiddleware.wrapStream = async (args) => { | |
| const startTime = Date.now(); | |
| try { | |
| const result = await middleware.wrapStream!(args); | |
| this.recordExecution(middleware.metadata.id, startTime, true); | |
| return result; | |
| } catch (error) { | |
| this.recordExecution( | |
| middleware.metadata.id, | |
| startTime, | |
| false, | |
| error as Error, | |
| ); | |
| throw error; | |
| } | |
| }; | |
| } | |
| return wrappedMiddleware; | |
| } | |
| // Create wrapper that tracks execution | |
| const wrappedMiddleware: LanguageModelV1Middleware = {}; | |
| if (middleware.transformParams) { | |
| wrappedMiddleware.transformParams = async (args) => { | |
| const startTime = Date.now(); | |
| try { | |
| const nextArgs = { | |
| ...args, | |
| params: { | |
| ...args.params, | |
| __neurolink: { | |
| ...(args.params?.__neurolink ?? {}), | |
| middlewareConfig: { | |
| ...(args.params?.__neurolink?.middlewareConfig ?? {}), | |
| [middleware.metadata.id]: mergedConfig, | |
| }, | |
| }, | |
| }, | |
| }; | |
| const result = await middleware.transformParams!(nextArgs); | |
| this.recordExecution(middleware.metadata.id, startTime, true); | |
| return result; | |
| } catch (error) { | |
| this.recordExecution( | |
| middleware.metadata.id, | |
| startTime, | |
| false, | |
| error as Error, | |
| ); | |
| throw error; | |
| } | |
| }; | |
| } | |
| if (middleware.wrapGenerate) { | |
| wrappedMiddleware.wrapGenerate = async (args) => { | |
| const startTime = Date.now(); | |
| try { | |
| const nextArgs = { | |
| ...args, | |
| params: { | |
| ...args.params, | |
| __neurolink: { | |
| ...(args.params?.__neurolink ?? {}), | |
| middlewareConfig: { | |
| ...(args.params?.__neurolink?.middlewareConfig ?? {}), | |
| [middleware.metadata.id]: mergedConfig, | |
| }, | |
| }, | |
| }, | |
| }; | |
| const result = await middleware.wrapGenerate!(nextArgs); | |
| this.recordExecution(middleware.metadata.id, startTime, true); | |
| return result; | |
| } catch (error) { | |
| this.recordExecution( | |
| middleware.metadata.id, | |
| startTime, | |
| false, | |
| error as Error, | |
| ); | |
| throw error; | |
| } | |
| }; | |
| } | |
| if (middleware.wrapStream) { | |
| wrappedMiddleware.wrapStream = async (args) => { | |
| const startTime = Date.now(); | |
| try { | |
| const nextArgs = { | |
| ...args, | |
| params: { | |
| ...args.params, | |
| __neurolink: { | |
| ...(args.params?.__neurolink ?? {}), | |
| middlewareConfig: { | |
| ...(args.params?.__neurolink?.middlewareConfig ?? {}), | |
| [middleware.metadata.id]: mergedConfig, | |
| }, | |
| }, | |
| }, | |
| }; | |
| const result = await middleware.wrapStream!(nextArgs); | |
| this.recordExecution(middleware.metadata.id, startTime, true); | |
| return result; | |
| } catch (error) { | |
| this.recordExecution( | |
| middleware.metadata.id, | |
| startTime, | |
| false, | |
| error as Error, | |
| ); | |
| throw error; | |
| } | |
| }; | |
| } | |
| return wrappedMiddleware; |
🤖 Prompt for AI Agents
In src/lib/middleware/registry.ts around lines 218 to 276, the wrapped
middleware calls need to receive the merged configuration in their args while
preserving existing metrics recording; update each wrapper (transformParams,
wrapGenerate, wrapStream) to compute the mergedConfig (existing registry +
middleware defaults) and call the original middleware with args where
args.params is augmented to include config: mergedConfig (e.g., newArgs = {
...args, params: { ...args.params, config: mergedConfig } }) before awaiting the
result, leaving the startTime/recordExecution try/catch logic unchanged.
| it("should track analytics in middleware", async () => { | ||
| const analyticsMiddleware = createAnalyticsMiddleware(); | ||
|
|
||
| // Mock the doGenerate function | ||
| const mockResult = { | ||
| text: "Hello, world!", | ||
| usage: { | ||
| promptTokens: 10, | ||
| completionTokens: 5, | ||
| }, | ||
| }; | ||
|
|
||
| const mockDoGenerate = vi.fn().mockResolvedValue(mockResult); | ||
|
|
||
| // Create a mock args object that satisfies TypeScript | ||
| const mockArgs = { | ||
| doGenerate: mockDoGenerate, | ||
| params: { prompt: "test" }, | ||
| // These are mocked to satisfy the type system | ||
| model: {} as unknown, | ||
| doStream: vi.fn().mockResolvedValue({ | ||
| stream: {} as unknown, | ||
| rawCall: { rawPrompt: {}, rawSettings: {} }, | ||
| }), | ||
| }; | ||
|
|
||
| // Use type assertion to bypass type checking for the test | ||
| const result = await (analyticsMiddleware.wrapGenerate as Function)( | ||
| mockArgs, | ||
| ); | ||
|
|
||
| expect(mockDoGenerate).toHaveBeenCalled(); | ||
| expect(result.text).toBe("Hello, world!"); | ||
| }); |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Add tests for wrapStream and error path to cover both middleware branches.
Current tests only cover wrapGenerate success. Add:
- a stream test to verify
rawResponse.neurolink.analyticsis set, and - an error test to ensure errors are logged and rethrown.
[tests]
Append the following tests:
@@
it("should track analytics in middleware", async () => {
@@
expect(result.text).toBe("Hello, world!");
+ // see above metadata assertions
});
+
+ it("should annotate stream analytics in rawResponse", async () => {
+ const analyticsMiddleware = createAnalyticsMiddleware();
+ const mockDoStream = vi.fn().mockResolvedValue({
+ stream: {} as unknown,
+ rawCall: { rawPrompt: {}, rawSettings: {} },
+ });
+ const result = await (analyticsMiddleware.wrapStream as Function)({
+ doStream: mockDoStream,
+ params: { prompt: "stream-test" },
+ model: {} as unknown,
+ });
+ expect(mockDoStream).toHaveBeenCalled();
+ expect(result.rawResponse?.neurolink?.analytics).toBeTruthy();
+ expect(result.rawResponse.neurolink.analytics.streamingMode).toBe(true);
+ });
+
+ it("should rethrow errors and log in wrapGenerate", async () => {
+ const analyticsMiddleware = createAnalyticsMiddleware();
+ const mockDoGenerate = vi.fn().mockRejectedValue(new Error("boom"));
+ await expect(
+ (analyticsMiddleware.wrapGenerate as Function)({
+ doGenerate: mockDoGenerate,
+ params: { prompt: "err" },
+ model: {} as unknown,
+ doStream: vi.fn(), // unused
+ }),
+ ).rejects.toThrow("boom");
+ expect(mockDoGenerate).toHaveBeenCalled();
+ });📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| it("should track analytics in middleware", async () => { | |
| const analyticsMiddleware = createAnalyticsMiddleware(); | |
| // Mock the doGenerate function | |
| const mockResult = { | |
| text: "Hello, world!", | |
| usage: { | |
| promptTokens: 10, | |
| completionTokens: 5, | |
| }, | |
| }; | |
| const mockDoGenerate = vi.fn().mockResolvedValue(mockResult); | |
| // Create a mock args object that satisfies TypeScript | |
| const mockArgs = { | |
| doGenerate: mockDoGenerate, | |
| params: { prompt: "test" }, | |
| // These are mocked to satisfy the type system | |
| model: {} as unknown, | |
| doStream: vi.fn().mockResolvedValue({ | |
| stream: {} as unknown, | |
| rawCall: { rawPrompt: {}, rawSettings: {} }, | |
| }), | |
| }; | |
| // Use type assertion to bypass type checking for the test | |
| const result = await (analyticsMiddleware.wrapGenerate as Function)( | |
| mockArgs, | |
| ); | |
| expect(mockDoGenerate).toHaveBeenCalled(); | |
| expect(result.text).toBe("Hello, world!"); | |
| }); | |
| it("should track analytics in middleware", async () => { | |
| const analyticsMiddleware = createAnalyticsMiddleware(); | |
| // Mock the doGenerate function | |
| const mockResult = { | |
| text: "Hello, world!", | |
| usage: { | |
| promptTokens: 10, | |
| completionTokens: 5, | |
| }, | |
| }; | |
| const mockDoGenerate = vi.fn().mockResolvedValue(mockResult); | |
| // Create a mock args object that satisfies TypeScript | |
| const mockArgs = { | |
| doGenerate: mockDoGenerate, | |
| params: { prompt: "test" }, | |
| // These are mocked to satisfy the type system | |
| model: {} as unknown, | |
| doStream: vi.fn().mockResolvedValue({ | |
| stream: {} as unknown, | |
| rawCall: { rawPrompt: {}, rawSettings: {} }, | |
| }), | |
| }; | |
| // Use type assertion to bypass type checking for the test | |
| const result = await (analyticsMiddleware.wrapGenerate as Function)( | |
| mockArgs, | |
| ); | |
| expect(mockDoGenerate).toHaveBeenCalled(); | |
| expect(result.text).toBe("Hello, world!"); | |
| // see above metadata assertions | |
| }); | |
| it("should annotate stream analytics in rawResponse", async () => { | |
| const analyticsMiddleware = createAnalyticsMiddleware(); | |
| const mockDoStream = vi.fn().mockResolvedValue({ | |
| stream: {} as unknown, | |
| rawCall: { rawPrompt: {}, rawSettings: {} }, | |
| }); | |
| const result = await (analyticsMiddleware.wrapStream as Function)({ | |
| doStream: mockDoStream, | |
| params: { prompt: "stream-test" }, | |
| model: {} as unknown, | |
| }); | |
| expect(mockDoStream).toHaveBeenCalled(); | |
| expect(result.rawResponse?.neurolink?.analytics).toBeTruthy(); | |
| expect(result.rawResponse.neurolink.analytics.streamingMode).toBe(true); | |
| }); | |
| it("should rethrow errors and log in wrapGenerate", async () => { | |
| const analyticsMiddleware = createAnalyticsMiddleware(); | |
| const mockDoGenerate = vi.fn().mockRejectedValue(new Error("boom")); | |
| await expect( | |
| (analyticsMiddleware.wrapGenerate as Function)({ | |
| doGenerate: mockDoGenerate, | |
| params: { prompt: "err" }, | |
| model: {} as unknown, | |
| doStream: vi.fn(), // unused | |
| }), | |
| ).rejects.toThrow("boom"); | |
| expect(mockDoGenerate).toHaveBeenCalled(); | |
| }); |
🤖 Prompt for AI Agents
In test/middleware-basic.test.ts around lines 42 to 75, current tests only
exercise wrapGenerate success; add two new tests: one for wrapStream that mocks
doStream to return a stream and a rawResponse object, calls
analyticsMiddleware.wrapStream and asserts that the returned rawResponse (or
emitted rawCall/rawResponse) has rawResponse.neurolink.analytics set as expected
and that the underlying doStream was invoked; and one for the error path that
mocks the wrapped function (doGenerate or doStream) to throw/reject, spies on
the logger.error mock, invokes the middleware wrapper, asserts the error is
rethrown (use expect(...).rejects.toThrow or try/catch with fail), and asserts
logger.error was called with the error; ensure mocks satisfy TypeScript types
and use vi.fn().mockResolvedValue / mockRejectedValue as appropriate.
| expect(mockDoGenerate).toHaveBeenCalled(); | ||
| expect(result.text).toBe("Hello, world!"); | ||
| }); |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Assert analytics metadata is attached to the result.
Analytics middleware injects experimental_providerMetadata.neurolink.analytics; verify it to catch regressions.
[tests]
expect(mockDoGenerate).toHaveBeenCalled();
expect(result.text).toBe("Hello, world!");
+ expect(
+ result.experimental_providerMetadata?.neurolink?.analytics,
+ ).toBeTruthy();
+ expect(
+ result.experimental_providerMetadata.neurolink.analytics.totalTokens,
+ ).toBe(15);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| expect(mockDoGenerate).toHaveBeenCalled(); | |
| expect(result.text).toBe("Hello, world!"); | |
| }); | |
| expect(mockDoGenerate).toHaveBeenCalled(); | |
| expect(result.text).toBe("Hello, world!"); | |
| expect( | |
| result.experimental_providerMetadata?.neurolink?.analytics, | |
| ).toBeTruthy(); | |
| expect( | |
| result.experimental_providerMetadata.neurolink.analytics.totalTokens, | |
| ).toBe(15); | |
| }); |
🤖 Prompt for AI Agents
In test/middleware-basic.test.ts around lines 73 to 75, the test currently only
checks that mockDoGenerate was called and the result text equals "Hello, world!"
but does not assert that the analytics metadata injected by the analytics
middleware is present; update the test to assert that
result.experimental_providerMetadata exists and includes neurolink.analytics
(e.g., check that result.experimental_providerMetadata.neurolink?.analytics is
defined and matches the expected shape or contains required fields), so the test
fails if the middleware stops attaching analytics metadata.
- This commit introduces a comprehensive guide for developing custom middleware in the NeuroLink platform. The guide covers the middleware architecture, key benefits, and the development process, enabling developers to extend and customize the platform's functionality.
The new documentation includes:
- An overview of the middleware architecture and its benefits
- A step-by-step guide to creating, registering, and applying custom middleware
- Best practices for middleware development, including error handling and testing
This guide will empower developers to build more modular, reusable, and extensible AI applications with NeuroLink.
2372e21 to
3a36901
Compare
|
🎉 This PR is included in version 7.25.0 🎉 The release is available on: Your semantic-release bot 📦🚀 |
2 participants
Description
This pull request introduces a comprehensive guide for developing custom middleware in the NeuroLink platform. The guide covers the middleware architecture, key benefits, and the development process, enabling developers to extend and customize the platform's functionality.
Type of Change
Related Issues
Changes Made
docs/CUSTOM-MIDDLEWARE-GUIDE.md, which provides a comprehensive guide to developing and implementing custom middleware in the NeuroLink platform.src/lib/middleware/builtin/analytics.tsandsrc/lib/middleware/factory.ts.AI Provider Impact
Component Impact
Testing
Test Environment
Performance Impact
Breaking Changes
Screenshots/Demo
Checklist
Additional Notes
Summary by CodeRabbit
New Features
Documentation
Tests