@mastra/rag 0.1.19-alpha.2 → 0.1.19-alpha.4

package/src/document/extractors/types.ts

@@ -0,0 +1,40 @@
+import { createOpenAI } from '@ai-sdk/openai';
+import type { MastraLanguageModel } from '@mastra/core/agent';
+import type {
+  KeywordExtractPrompt,
+  QuestionExtractPrompt,
+  SummaryPrompt,
+  TitleExtractorPrompt,
+  TitleCombinePrompt,
+} from 'llamaindex';
+
+export type KeywordExtractArgs = {
+  llm?: MastraLanguageModel;
+  keywords?: number;
+  promptTemplate?: KeywordExtractPrompt['template'];
+};
+
+export type QuestionAnswerExtractArgs = {
+  llm?: MastraLanguageModel;
+  questions?: number;
+  promptTemplate?: QuestionExtractPrompt['template'];
+  embeddingOnly?: boolean;
+};
+
+export type SummaryExtractArgs = {
+  llm?: MastraLanguageModel;
+  summaries?: string[];
+  promptTemplate?: SummaryPrompt['template'];
+};
+
+export type TitleExtractorsArgs = {
+  llm?: MastraLanguageModel;
+  nodes?: number;
+  nodeTemplate?: TitleExtractorPrompt['template'];
+  combineTemplate?: TitleCombinePrompt['template'];
+};
+
+export const STRIP_REGEX = /(\r\n|\n|\r)/gm;
+
+const openai = createOpenAI({ apiKey: process.env.OPENAI_API_KEY });
+export const baseLLM: MastraLanguageModel = openai('gpt-4o');

package/src/document/types.ts

@@ -1,12 +1,10 @@
 import type { TiktokenEncoding, TiktokenModel } from 'js-tiktoken';
 import type {
-  LLM,
-  TitleCombinePrompt,
-  TitleExtractorPrompt,
-  SummaryPrompt,
-  QuestionExtractPrompt,
-  KeywordExtractPrompt,
-} from 'llamaindex';
+  TitleExtractorsArgs,
+  SummaryExtractArgs,
+  QuestionAnswerExtractArgs,
+  KeywordExtractArgs,
+} from './extractors';
 
 export enum Language {
   CPP = 'cpp',
@@ -69,32 +67,6 @@ export type ChunkOptions = {
   stripHeaders?: boolean;
 };
 
-export type TitleExtractorsArgs = {
-  llm?: LLM;
-  nodes?: number;
-  nodeTemplate?: TitleExtractorPrompt['template'];
-  combineTemplate?: TitleCombinePrompt['template'];
-};
-
-export type SummaryExtractArgs = {
-  llm?: LLM;
-  summaries?: string[];
-  promptTemplate?: SummaryPrompt['template'];
-};
-
-export type QuestionAnswerExtractArgs = {
-  llm?: LLM;
-  questions?: number;
-  promptTemplate?: QuestionExtractPrompt['template'];
-  embeddingOnly?: boolean;
-};
-
-export type KeywordExtractArgs = {
-  llm?: LLM;
-  keywords?: number;
-  promptTemplate?: KeywordExtractPrompt['template'];
-};
-
 export type ChunkStrategy = 'recursive' | 'character' | 'token' | 'markdown' | 'html' | 'json' | 'latex';
 
 export interface ChunkParams extends ChunkOptions {