@jmndao/mongoose-ai
Version:
AI-powered Mongoose plugin for intelligent document processing with auto-summarization, semantic search, MongoDB Vector Search, and function calling
703 lines (676 loc) • 20.6 kB
TypeScript
import { Document, Model, Schema } from 'mongoose';
/**
* Core types and enums
*/
type AIModel = "summary" | "embedding";
type AIProvider = "openai" | "anthropic" | "ollama";
type LogLevel = "debug" | "info" | "warn" | "error";
/**
* AI provider credentials
*/
interface AICredentials {
apiKey: string;
organizationId?: string;
}
/**
* Function calling types
*/
/**
* Function parameter definition
*/
interface FunctionParameter {
type: "string" | "number" | "boolean" | "array" | "object";
description: string;
enum?: string[] | number[];
items?: FunctionParameter;
properties?: Record<string, FunctionParameter>;
required?: boolean;
}
/**
* Function definition
*/
interface AIFunction {
name: string;
description: string;
parameters: Record<string, FunctionParameter>;
handler: (args: any, document: Document) => Promise<void> | void;
}
/**
* Function execution result
*/
interface FunctionResult {
name: string;
success: boolean;
result?: any;
error?: string;
executedAt: Date;
}
/**
* Quick function builders
*/
declare const QuickFunctions: {
updateField: (fieldName: string, allowedValues?: string[]) => AIFunction;
scoreField: (fieldName: string, min?: number, max?: number) => AIFunction;
manageTags: (fieldName?: string) => AIFunction;
};
/**
* Create a custom function
*/
declare function createFunction(name: string, description: string, parameters: Record<string, FunctionParameter>, handler: (args: any, document: Document) => Promise<void> | void): AIFunction;
/**
* Configuration types
*/
/**
* Advanced configuration options
*/
interface AIAdvancedOptions {
/** Maximum retries for failed API calls (default: 2) */
maxRetries?: number;
/** Timeout for API calls in milliseconds (default: 30000) */
timeout?: number;
/** Skip AI processing on document updates (default: false) */
skipOnUpdate?: boolean;
/** Force regeneration even if AI content exists (default: false) */
forceRegenerate?: boolean;
/** Log level for debugging (default: 'warn') */
logLevel?: LogLevel;
/** Continue saving document if AI processing fails (default: true) */
continueOnError?: boolean;
/** Enable function calling (default: false) */
enableFunctions?: boolean;
}
/**
* OpenAI specific model configuration
*/
interface OpenAIModelConfig {
/** Chat model for summaries (default: 'gpt-3.5-turbo') */
chatModel?: string;
/** Embedding model (default: 'text-embedding-3-small') */
embeddingModel?: string;
/** Max tokens for summaries (default: 200) */
maxTokens?: number;
/** Temperature for text generation (default: 0.3) */
temperature?: number;
}
/**
* Anthropic specific model configuration
*/
interface AnthropicModelConfig {
/** Chat model for summaries (default: 'claude-3-haiku-20240307') */
chatModel?: string;
/** Max tokens for summaries (default: 200) */
maxTokens?: number;
/** Temperature for text generation (default: 0.3) */
temperature?: number;
}
/**
* Ollama specific model configuration
*/
interface OllamaModelConfig {
/** Chat model for summaries (default: 'llama3.2') */
chatModel?: string;
/** Embedding model (default: 'nomic-embed-text') */
embeddingModel?: string;
/** Max tokens for summaries (default: 200) */
maxTokens?: number;
/** Temperature for text generation (default: 0.3) */
temperature?: number;
/** Ollama server endpoint (default: 'http://localhost:11434') */
endpoint?: string;
}
/**
* Vector search configuration
*/
interface VectorSearchConfig {
/** Enable MongoDB Vector Search (default: auto-detect) */
enabled?: boolean;
/** Vector search index name (default: 'vector_index') */
indexName?: string;
/** Auto-create vector index if not exists (default: true) */
autoCreateIndex?: boolean;
/** Vector index similarity metric (default: 'cosine') */
similarity?: "cosine" | "euclidean" | "dotProduct";
}
/**
* Main AI configuration
*/
interface AIConfig {
model: AIModel;
provider: AIProvider;
field: string;
credentials: AICredentials;
prompt?: string;
advanced?: AIAdvancedOptions;
modelConfig?: OpenAIModelConfig | AnthropicModelConfig | OllamaModelConfig;
vectorSearch?: VectorSearchConfig;
/** Fields to include in AI processing */
includeFields?: string[];
/** Fields to exclude from AI processing */
excludeFields?: string[];
/** Functions to make available to AI */
functions?: AIFunction[];
}
/**
* Plugin options
*/
interface AIPluginOptions {
ai: AIConfig;
}
/**
* Result and response types
*/
/**
* AI-generated summary result
*/
interface SummaryResult {
summary: string;
generatedAt: Date;
model: string;
tokenCount?: number;
processingTime?: number;
functionResults?: FunctionResult[];
}
/**
* AI-generated embedding result
*/
interface EmbeddingResult {
embedding: number[];
generatedAt: Date;
model: string;
dimensions: number;
processingTime?: number;
functionResults?: FunctionResult[];
}
/**
* Processing statistics
*/
interface AIProcessingStats {
totalProcessed: number;
successful: number;
failed: number;
averageProcessingTime: number;
totalTokensUsed?: number;
}
/**
* AI error information
*/
interface AIError {
message: string;
code: string;
originalError?: any;
timestamp: Date;
retryCount?: number;
}
/**
* Search and similarity types
*/
/**
* Semantic search result
*/
interface SearchResult<T = any> {
document: T;
similarity: number;
metadata?: {
field: string;
distance: number;
};
}
/**
* Semantic search options
*/
interface SemanticSearchOptions {
/** Maximum number of results (default: 10) */
limit?: number;
/** Minimum similarity threshold (default: 0.7) */
threshold?: number;
/** Include similarity scores (default: true) */
includeScore?: boolean;
/** Additional MongoDB query filters */
filter?: Record<string, any>;
/** Use MongoDB Vector Search instead of in-memory search (default: auto-detect) */
useVectorSearch?: boolean;
/** Vector search index name (default: 'vector_index') */
indexName?: string;
/** Number of candidates for vector search (default: limit * 10) */
numCandidates?: number;
}
/**
* Type guard for semantic search results
*/
declare function isSearchResult<T>(obj: any): obj is SearchResult<T>;
/**
* Model and document extension types
*/
/**
* Extended document interface with AI methods
*/
interface AIDocumentMethods {
getAIContent(): SummaryResult | EmbeddingResult | null;
regenerateAI(): Promise<void>;
calculateSimilarity?(other: any): number;
}
/**
* Extended model interface with semantic search methods
*/
interface AIModelStatics<T = any> {
semanticSearch(query: string, options?: SemanticSearchOptions): Promise<SearchResult<T>[]>;
findSimilar(document: T, options?: SemanticSearchOptions): Promise<SearchResult<T>[]>;
getAIStats?(): AIProcessingStats;
resetAIStats?(): void;
}
/**
* Combined AI document type
*/
type AIDocument<T = any> = Document & T & AIDocumentMethods;
/**
* Combined AI model type
*/
type AIModelType<T = any> = Model<T> & AIModelStatics<T>;
/**
* Type helper to add AI methods to existing model type
*/
type WithAI<TModel extends Model<any>> = TModel & {
semanticSearch(query: string, options?: SemanticSearchOptions): Promise<SearchResult<any>[]>;
findSimilar(document: any, options?: SemanticSearchOptions): Promise<SearchResult<any>[]>;
};
/**
* Type helper to add AI methods to existing document type
*/
type WithAIDocument<TDoc extends Document> = TDoc & {
getAIContent(): SummaryResult | EmbeddingResult | null;
regenerateAI(): Promise<void>;
calculateSimilarity?(other: any): number;
};
/**
* Check if a model has AI methods
*/
declare function hasAIMethods<T>(model: Model<T>): model is Model<T> & AIModelStatics<T>;
/**
* Check if a document has AI methods
*/
declare function hasAIDocumentMethods<T extends Document>(doc: T): doc is T & AIDocumentMethods;
/**
* Configuration helpers and utilities
*/
/**
* Default configurations
*/
declare const DEFAULT_CONFIG: {
advanced: {
maxRetries: number;
timeout: number;
skipOnUpdate: boolean;
forceRegenerate: boolean;
logLevel: "warn";
continueOnError: boolean;
enableFunctions: boolean;
};
openai: {
chatModel: string;
embeddingModel: string;
maxTokens: number;
temperature: number;
};
anthropic: {
chatModel: string;
maxTokens: number;
temperature: number;
};
ollama: {
chatModel: string;
embeddingModel: string;
maxTokens: number;
temperature: number;
endpoint: string;
};
vectorSearch: {
enabled: boolean;
indexName: string;
autoCreateIndex: boolean;
similarity: "cosine";
};
};
/**
* Validate API key format
*/
declare function validateApiKey(apiKey: string, provider: AIProvider): boolean;
/**
* Create AI configuration with all options including vector search
*/
declare function createAdvancedAIConfig(options: {
apiKey: string;
provider: AIProvider;
model: AIModel;
field: string;
prompt?: string;
advanced?: Partial<AIAdvancedOptions>;
modelConfig?: Partial<OpenAIModelConfig | AnthropicModelConfig | OllamaModelConfig>;
vectorSearch?: Partial<VectorSearchConfig>;
includeFields?: string[];
excludeFields?: string[];
functions?: any[];
}): AIConfig;
/**
* Create basic AI configuration (maintains v1.0.x compatibility)
*/
declare function createAIConfig(options: {
apiKey: string;
model: AIModel;
field: string;
prompt?: string;
advanced?: Partial<AIAdvancedOptions>;
modelConfig?: Partial<OpenAIModelConfig>;
includeFields?: string[];
excludeFields?: string[];
}): AIConfig;
/**
* Create Ollama AI configuration helper
*/
declare function createOllamaConfig(options: {
model: AIModel;
field: string;
endpoint?: string;
chatModel?: string;
embeddingModel?: string;
prompt?: string;
advanced?: Partial<AIAdvancedOptions>;
vectorSearch?: Partial<VectorSearchConfig>;
includeFields?: string[];
excludeFields?: string[];
functions?: any[];
}): AIConfig;
/**
* Estimate token count for text (rough estimation)
*/
declare function estimateTokenCount(text: string): number;
/**
* Estimate cost for AI operations
*/
declare function estimateCost(tokenCount: number, model: string, provider?: AIProvider): number;
/**
* Check if environment is properly configured
*/
declare function checkEnvironment(): {
isValid: boolean;
missing: string[];
warnings: string[];
};
/**
* Main mongoose-ai plugin with MongoDB Vector Search and Ollama support
*/
/**
* Main plugin function
*/
declare function aiPlugin(schema: Schema, options: AIPluginOptions): void;
/**
* Base provider interface for AI operations
*/
/**
* Base AI provider interface
*/
interface AIProviderInterface {
/**
* Generate summary for document
*/
summarize(document: Record<string, any>, customPrompt?: string, functions?: AIFunction[]): Promise<SummaryResult>;
/**
* Generate embedding for text
*/
generateEmbedding(text: string): Promise<EmbeddingResult>;
/**
* Get provider information
*/
getProviderInfo(): {
name: string;
version: string;
models: any;
advanced: any;
};
}
/**
* Base provider class with common functionality
*/
declare abstract class BaseProvider implements AIProviderInterface {
protected readonly advanced: Required<AIAdvancedOptions>;
constructor(credentials: AICredentials, advancedOptions?: AIAdvancedOptions);
abstract summarize(document: Record<string, any>, customPrompt?: string, functions?: AIFunction[]): Promise<SummaryResult>;
abstract generateEmbedding(text: string): Promise<EmbeddingResult>;
abstract getProviderInfo(): {
name: string;
version: string;
models: any;
advanced: any;
};
/**
* Execute functions with error handling
*/
protected executeFunctions(functions: AIFunction[], functionCalls: {
name: string;
arguments: any;
}[], document: Document): Promise<FunctionResult[]>;
/**
* Convert document to clean text
*/
protected prepareText(document: Record<string, any>): string;
/**
* Truncate text to specified length
*/
protected truncateText(text: string, maxLength: number): string;
/**
* Validate credentials
*/
protected abstract validateCredentials(credentials: AICredentials): void;
/**
* Extract error message
*/
protected getErrorMessage(error: any): string;
/**
* Log messages based on level
*/
protected log(level: LogLevel, message: string, error?: any): void;
}
/**
* Anthropic provider for text processing with function calling
*/
declare class AnthropicProvider extends BaseProvider {
private readonly config;
private readonly apiKey;
constructor(credentials: AICredentials, modelConfig?: AnthropicModelConfig, advancedOptions?: AIAdvancedOptions);
/**
* Generate summary for document with optional function calling
*/
summarize(document: Record<string, any>, customPrompt?: string, functions?: AIFunction[]): Promise<SummaryResult>;
/**
* Generate embedding for text - Note: Anthropic doesn't provide embeddings
*/
generateEmbedding(text: string): Promise<EmbeddingResult>;
/**
* Make API request to Anthropic
*/
private makeAnthropicRequest;
/**
* Prepare function tools for Anthropic API
*/
private prepareFunctionTools;
/**
* Validate credentials
*/
protected validateCredentials(credentials: AICredentials): void;
/**
* Get provider information
*/
getProviderInfo(): {
name: string;
version: string;
models: Required<AnthropicModelConfig>;
advanced: Required<AIAdvancedOptions>;
};
}
/**
* OpenAI provider for text processing and embeddings with function calling
*/
declare class OpenAIProvider extends BaseProvider {
private readonly client;
private readonly config;
constructor(credentials: AICredentials, modelConfig?: OpenAIModelConfig, advancedOptions?: AIAdvancedOptions);
/**
* Generate summary for document with optional function calling
*/
summarize(document: Record<string, any>, customPrompt?: string, functions?: AIFunction[]): Promise<SummaryResult>;
/**
* Generate embedding for text
*/
generateEmbedding(text: string): Promise<EmbeddingResult>;
/**
* Prepare function tools for OpenAI API
*/
private prepareFunctionTools;
/**
* Validate credentials
*/
protected validateCredentials(credentials: AICredentials): void;
/**
* Get provider information
*/
getProviderInfo(): {
name: string;
version: string;
models: Required<OpenAIModelConfig>;
advanced: Required<AIAdvancedOptions>;
};
}
/**
* Ollama provider for local LLM processing
*/
declare class OllamaProvider extends BaseProvider {
private readonly config;
private readonly endpoint;
constructor(credentials: AICredentials, modelConfig?: OllamaModelConfig, advancedOptions?: AIAdvancedOptions);
/**
* Generate summary for document with optional function calling
*/
summarize(document: Record<string, any>, customPrompt?: string, functions?: AIFunction[]): Promise<SummaryResult>;
/**
* Generate embedding for text
*/
generateEmbedding(text: string): Promise<EmbeddingResult>;
/**
* Make API request to Ollama
*/
private makeOllamaRequest;
/**
* Prepare function calling instructions for Ollama
*/
private prepareFunctionInstructions;
/**
* Parse function calls from Ollama response
*/
private parseFunctionCalls;
/**
* Parse simple function arguments
*/
private parseSimpleArgs;
/**
* Estimate token count (rough approximation)
*/
private estimateTokenCount;
/**
* Validate credentials
*/
protected validateCredentials(credentials: AICredentials): void;
/**
* Get provider information
*/
getProviderInfo(): {
name: string;
version: string;
models: Required<OllamaModelConfig>;
advanced: Required<AIAdvancedOptions>;
endpoint: string;
};
}
/**
* MongoDB Vector Search utilities
*/
/**
* Check if MongoDB Vector Search is available
*/
declare function detectVectorSearchSupport(model: Model<any>): Promise<boolean>;
/**
* Create vector search index for embeddings
*/
declare function createVectorIndex(model: Model<any>, embeddingField: string, dimensions: number, config: VectorSearchConfig): Promise<void>;
/**
* Calculate cosine similarity between two vectors (fallback for in-memory search)
*/
declare function cosineSimilarity(a: number[], b: number[]): number;
/**
* Package version
*/
declare const VERSION = "1.4.0";
/**
* Supported models and providers
*/
declare const SUPPORTED_MODELS: readonly ["summary", "embedding"];
declare const SUPPORTED_PROVIDERS: readonly ["openai", "anthropic", "ollama"];
/**
* Default export with core functionality (for backward compatibility)
*/
declare const mongooseAI: {
aiPlugin: typeof aiPlugin;
OpenAIProvider: typeof OpenAIProvider;
AnthropicProvider: typeof AnthropicProvider;
OllamaProvider: typeof OllamaProvider;
QuickFunctions: {
updateField: (fieldName: string, allowedValues?: string[]) => AIFunction;
scoreField: (fieldName: string, min?: number, max?: number) => AIFunction;
manageTags: (fieldName?: string) => AIFunction;
};
createFunction: typeof createFunction;
createAIConfig: typeof createAIConfig;
createAdvancedAIConfig: typeof createAdvancedAIConfig;
createOllamaConfig: typeof createOllamaConfig;
validateApiKey: typeof validateApiKey;
estimateTokenCount: typeof estimateTokenCount;
estimateCost: typeof estimateCost;
checkEnvironment: typeof checkEnvironment;
VERSION: string;
SUPPORTED_MODELS: readonly ["summary", "embedding"];
SUPPORTED_PROVIDERS: readonly ["openai", "anthropic", "ollama"];
DEFAULT_CONFIG: {
advanced: {
maxRetries: number;
timeout: number;
skipOnUpdate: boolean;
forceRegenerate: boolean;
logLevel: "warn";
continueOnError: boolean;
enableFunctions: boolean;
};
openai: {
chatModel: string;
embeddingModel: string;
maxTokens: number;
temperature: number;
};
anthropic: {
chatModel: string;
maxTokens: number;
temperature: number;
};
ollama: {
chatModel: string;
embeddingModel: string;
maxTokens: number;
temperature: number;
endpoint: string;
};
vectorSearch: {
enabled: boolean;
indexName: string;
autoCreateIndex: boolean;
similarity: "cosine";
};
};
};
export { type AIAdvancedOptions, type AIConfig, type AICredentials, type AIDocument, type AIDocumentMethods, type AIError, type AIFunction, type AIModel, type AIModelStatics, type AIModelType, type AIPluginOptions, type AIProcessingStats, type AIProvider, type AnthropicModelConfig, AnthropicProvider, DEFAULT_CONFIG, type EmbeddingResult, type FunctionParameter, type FunctionResult, type LogLevel, type OllamaModelConfig, OllamaProvider, type OpenAIModelConfig, OpenAIProvider, QuickFunctions, SUPPORTED_MODELS, SUPPORTED_PROVIDERS, type SearchResult, type SemanticSearchOptions, type SummaryResult, VERSION, type VectorSearchConfig, type WithAI, type WithAIDocument, aiPlugin, checkEnvironment, cosineSimilarity, createAIConfig, createAdvancedAIConfig, createFunction, createOllamaConfig, createVectorIndex, mongooseAI as default, detectVectorSearchSupport, estimateCost, estimateTokenCount, hasAIDocumentMethods, hasAIMethods, isSearchResult, validateApiKey };