@ixo/common 1.1.0

package/src/ai/tools/retriever-tool/retriever-tool.test.ts

@@ -0,0 +1,156 @@
+import { type VectorDBDataStore } from '@ixo/data-store';
+import { Document } from '@langchain/core/documents';
+import { FakeChatModel } from '@langchain/core/utils/testing';
+import { retrieverToolFactory } from './retriever-tool.js';
+
+const fakeModel = new FakeChatModel({});
+
+// Mock the logger to prevent console output during tests
+jest.mock('@ixo/logger', () => ({
+  Logger: {
+    error: jest.fn(),
+  },
+}));
+
+// Mock the OpenAI model
+jest.mock('../../models/openai', () => ({
+  getChatOpenAiModel: () => fakeModel,
+}));
+
+// Mock the doc relevance checker
+jest.mock('../../utils/doc-relevance-checker', () => ({
+  __esModule: true,
+  default: jest.fn().mockResolvedValue(true),
+}));
+const fakeDocRelevanceChecker = jest.fn().mockResolvedValue(true);
+
+jest
+  .spyOn(require('../../utils/doc-relevance-checker'), 'default')
+  .mockImplementation(fakeDocRelevanceChecker);
+
+describe('RetrieverTool', () => {
+  let mockStore: jest.Mocked<VectorDBDataStore>;
+  let mockModel: FakeChatModel = new FakeChatModel({});
+
+  beforeEach(() => {
+    mockStore = {
+      queryWithSimilarity: jest.fn(),
+    } as unknown as jest.Mocked<VectorDBDataStore>;
+  });
+
+  it('should retrieve documents with default settings', async () => {
+    const mockDocs = [
+      { id: '1', content: 'test content 1', metadata: { source: 'test1' } },
+      { id: '2', content: 'test content 2', metadata: { source: 'test2' } },
+    ];
+
+    mockStore.queryWithSimilarity.mockResolvedValue(mockDocs);
+
+    const tool = retrieverToolFactory({ store: mockStore });
+    const result = (await tool.invoke({ query: 'test query' })) as Document[];
+
+    expect(mockStore.queryWithSimilarity).toHaveBeenCalledWith('test query', {
+      similarityThreshold: 0.3,
+      filters: undefined,
+    });
+
+    expect(result).toHaveLength(2);
+    expect(result[0]).toBeInstanceOf(Document);
+    expect(result[0]?.pageContent).toBe('test content 1');
+    expect(result[0]?.metadata).toEqual({ source: 'test1' });
+
+    expect(fakeDocRelevanceChecker).toHaveBeenCalledWith({
+      doc: result[0],
+      query: 'test query',
+      model: fakeModel,
+    });
+  });
+
+  it('should store results in provided map', async () => {
+    const mockDocs = [
+      { id: '1', content: 'test content', metadata: { source: 'test' } },
+    ];
+    mockStore.queryWithSimilarity.mockResolvedValue(mockDocs);
+
+    const resultsMap = new Map();
+    const requestId = 'test-request-id';
+
+    const tool = retrieverToolFactory({
+      store: mockStore,
+      map: resultsMap,
+      requestId,
+    });
+
+    await tool.invoke({ query: 'test query' });
+
+    expect(resultsMap.get(requestId)).toEqual([{ source: 'test' }]);
+  });
+
+  it('should return undefined when no documents found', async () => {
+    mockStore.queryWithSimilarity.mockResolvedValue([]);
+
+    const tool = retrieverToolFactory({ store: mockStore });
+    const result = await tool.invoke({ query: 'test query' });
+
+    expect(result).toBeUndefined();
+  });
+
+  it('should check document relevance when similarity threshold is higher', async () => {
+    const mockDocs = [
+      { id: '1', content: 'test content', metadata: { source: 'test' } },
+    ];
+    mockStore.queryWithSimilarity.mockResolvedValue(mockDocs);
+
+    const tool = retrieverToolFactory({
+      store: mockStore,
+      similarThreshold: 0.5,
+      model: mockModel,
+    });
+
+    const result = await tool.invoke({ query: 'test query' });
+
+    expect(result).toHaveLength(1);
+    expect(result[0]).toBeInstanceOf(Document);
+  });
+
+  it('should handle errors during relevance checking', async () => {
+    const mockDocs = [
+      { id: '1', content: 'test content', metadata: { source: 'test' } },
+    ];
+    mockStore.queryWithSimilarity.mockResolvedValue(mockDocs);
+
+    // Mock the relevance checker to throw an error
+    jest
+      .requireMock('../../utils/doc-relevance-checker')
+      .default.mockRejectedValueOnce(new Error('Test error'));
+
+    const tool = retrieverToolFactory({
+      store: mockStore,
+      similarThreshold: 0.5,
+      model: mockModel,
+    });
+
+    const result = await tool.invoke({ query: 'test query' });
+
+    expect(result).toHaveLength(0);
+  });
+
+  it('should apply filters when provided', async () => {
+    const mockDocs = [
+      { id: '1', content: 'test content', metadata: { source: 'test' } },
+    ];
+    mockStore.queryWithSimilarity.mockResolvedValue(mockDocs);
+    const filters = { category: 'test' };
+    const tool = retrieverToolFactory({
+      store: mockStore,
+      filters,
+    });
+
+    await tool.invoke({ query: 'test query' });
+
+    expect(mockStore.queryWithSimilarity).toHaveBeenCalledWith('test query', {
+      similarityThreshold: 0.3,
+      filters,
+    });
+  });
+});

package/src/ai/tools/retriever-tool/retriever-tool.ts

@@ -0,0 +1,107 @@
+import { type VectorDBDataStore } from '@ixo/data-store';
+import { Logger } from '@ixo/logger';
+import { Document } from '@langchain/core/documents';
+import { type BaseChatModel } from '@langchain/core/language_models/chat_models';
+import { tool, type DynamicStructuredTool } from '@langchain/core/tools';
+import z from 'zod';
+import { getChatOpenAiModel } from '../../models/openai.js';
+import checkDocRelevance from '../../utils/doc-relevance-checker.js';
+
+type RetrieverToolFactoryArgs<
+  Filters extends Record<string, unknown> = Record<string, unknown>,
+  MapType extends Map<string, unknown> = Map<string, unknown>,
+> = {
+  model?: BaseChatModel;
+  filters?: Filters;
+  similarThreshold?: number;
+  store: VectorDBDataStore;
+  map?: MapType;
+  requestId?: string;
+};
+
+const schema = z.object({
+  query: z.string(),
+});
+
+export const retrieverToolFactory = ({
+  model = getChatOpenAiModel(),
+  filters,
+  similarThreshold = 0.3,
+  map,
+  requestId,
+  store,
+}: RetrieverToolFactoryArgs): DynamicStructuredTool<typeof schema> =>
+  tool(
+    async ({ query }) => {
+      const docs = await store.queryWithSimilarity(query, {
+        similarityThreshold: similarThreshold,
+        filters,
+      });
+
+      // save the results to the map if the parent class wan to access the docs used by the Agent
+      if (map) {
+        map.set(
+          requestId ?? query,
+          docs.map((doc) => doc.metadata),
+        );
+      }
+
+      if (docs.length === 0) {
+        return undefined;
+      }
+
+      if (similarThreshold >= 0.3) {
+        // Filter the docs based on relevance using the AI
+        const relevantDocs = await Promise.all(
+          docs.map(async (doc) => {
+            try {
+              return (
+                await checkDocRelevance({
+                  doc: new Document({
+                    pageContent: doc.content,
+                    id: doc.id,
+                    metadata: doc.metadata,
+                  }),
+                  query,
+                  model,
+                })
+              ).valueOf()
+                ? doc
+                : null;
+            } catch (error) {
+              Logger.error(`Error checking relevance of document:`, error);
+              return null;
+            }
+          }),
+        );
+
+        return relevantDocs.reduce<Document[]>((acc, doc) => {
+          if (doc) {
+            acc.push(
+              new Document({
+                pageContent: doc.content,
+                id: doc.id,
+                metadata: doc.metadata,
+              }),
+            );
+          }
+          return acc;
+        }, []);
+      }
+
+      return docs.map(
+        (doc) =>
+          new Document({
+            pageContent: doc.content,
+            id: doc.id,
+            metadata: doc.metadata,
+          }),
+      );
+    },
+    {
+      name: 'retrieverTool',
+      description:
+        'Retrieves documents from the knowledge base based on the query',
+      schema,
+    },
+  );

package/src/ai/tools/scrape-web-page.ts

@@ -0,0 +1,75 @@
+import { PlaywrightWebBaseLoader } from '@langchain/community/document_loaders/web/playwright';
+import { HtmlToTextTransformer } from '@langchain/community/document_transformers/html_to_text';
+import { Document } from '@langchain/core/documents';
+import { PromptTemplate } from '@langchain/core/prompts';
+import { ChatOpenAI } from '@langchain/openai';
+import { createAgent } from 'langchain';
+
+export const scrapeWebPage = async (url: string) => {
+  const content = await PlaywrightWebBaseLoader._scrape(url, {
+    launchOptions: {
+      headless: true,
+    },
+    gotoOptions: {
+      waitUntil: 'domcontentloaded',
+    },
+  });
+  const transformer = new HtmlToTextTransformer({});
+  const docs = await transformer.transformDocuments([
+    new Document({
+      pageContent: content,
+      metadata: {
+        source: url,
+      },
+    }),
+  ]);
+  return new Document({
+    pageContent: docs.map((doc) => doc.pageContent).join('\n'),
+    metadata: {
+      source: url,
+    },
+  });
+};
+const prompt = PromptTemplate.fromTemplate(
+  `
+  You are a helpful assistant that summarizes web pages.
+
+  #Guidelines
+  - Add key points at the beginning of the Document.
+  - Add the summary based on the key points.
+
+  #Rules
+  - Ignore the html tags and only focus on the content.
+  - Read the content of the page and summarize.
+  - DO NOT HALLUCINATE or make up information.
+  - DO NOT ADD ANYTHING ELSE OTHER THAN THE SUMMARY.
+
+  #Document
+  {context}
+  `,
+);
+
+export const scrapeAndSummarizeWebPage = async (
+  url: string,
+  llm: ChatOpenAI,
+) => {
+  const doc = await scrapeWebPage(url);
+
+  const agent = createAgent({
+    model: llm,
+    tools: [],
+  });
+
+  const result = await agent.invoke({
+    messages: [
+      {
+        role: 'user',
+        content: await prompt.format({
+          context: doc.pageContent,
+        }),
+      },
+    ],
+  });
+
+  return result;
+};

package/src/ai/tools/web-search-tool.ts

@@ -0,0 +1,38 @@
+import { tool } from '@langchain/core/tools';
+import { tavily } from '@tavily/core';
+import z from 'zod';
+import { jsonToYaml } from '../utils/json-to-yaml.js';
+
+const webSearchTool = tool(
+  async ({ input }: { input: string }) => {
+    const tvly = tavily({ apiKey: process.env.TAVILY_API_KEY });
+    const searchResults = await tvly.search(input, {
+      maxResults: 3,
+      includeAnswer: true,
+    });
+
+    const res = searchResults.results.map((result) =>
+      jsonToYaml({
+        title: result.title,
+        url: result.url,
+        content: result.content,
+        publishedAt: result.publishedDate,
+      }),
+    );
+    return jsonToYaml({
+      query: searchResults.query,
+      summary: searchResults.answer,
+      results: res,
+    }).replace(/[{}]/g, '');
+  },
+  {
+    name: 'web_search_tool',
+    description:
+      'A search engine optimized for comprehensive, accurate, and trusted results. Useful for when you need to answer questions about current events. Input should be a search query.',
+    schema: z.object({
+      input: z.string('The query to search for'),
+    }),
+  },
+);
+
+export { webSearchTool };

package/src/ai/types.ts

@@ -0,0 +1,6 @@
+export type EnsureKeys<
+  T extends Record<string, unknown>,
+  K extends (keyof T)[],
+> = {
+  [P in K[number]]: T[P]; // At least one key 'K' from 'T'
+} & Partial<T>; // Other keys of 'T' are optional

package/src/ai/utils/__tests__/chunk-arr.test.ts

@@ -0,0 +1,46 @@
+import { chunkArr } from '../chunk-arr.js';
+
+describe('chunkArr', () => {
+  it('should split array into chunks of specified size', () => {
+    const array = [1, 2, 3, 4, 5, 6, 7];
+    const result = chunkArr(array, 3);
+    expect(result).toEqual([[1, 2, 3], [4, 5, 6], [7]]);
+  });
+
+  it('should handle empty arrays', () => {
+    const array: number[] = [];
+    const result = chunkArr(array, 2);
+    expect(result).toEqual([]);
+  });
+
+  it('should handle chunk size equal to array length', () => {
+    const array = [1, 2, 3];
+    const result = chunkArr(array, 3);
+    expect(result).toEqual([[1, 2, 3]]);
+  });
+
+  it('should handle chunk size larger than array length', () => {
+    const array = [1, 2];
+    const result = chunkArr(array, 3);
+    expect(result).toEqual([[1, 2]]);
+  });
+
+  it('should throw error for chunk size less than or equal to 0', () => {
+    const array = [1, 2, 3];
+    expect(() => chunkArr(array, 0)).toThrow(
+      'Chunk size must be greater than 0',
+    );
+    expect(() => chunkArr(array, -1)).toThrow(
+      'Chunk size must be greater than 0',
+    );
+  });
+
+  it('should work with different types', () => {
+    const array = ['a', 'b', 'c', 'd'];
+    const result = chunkArr(array, 2);
+    expect(result).toEqual([
+      ['a', 'b'],
+      ['c', 'd'],
+    ]);
+  });
+});

package/src/ai/utils/__tests__/doc-relevance-checker.test.ts

@@ -0,0 +1,90 @@
+import { Document } from '@langchain/core/documents';
+import { type BaseChatModel } from '@langchain/core/language_models/chat_models';
+import { ChatPromptTemplate } from '@langchain/core/prompts';
+import { FakeChatModel } from '@langchain/core/utils/testing';
+import checkDocRelevance from '../doc-relevance-checker.js';
+
+jest.spyOn(ChatPromptTemplate, 'fromTemplate');
+
+describe('checkDocRelevance', () => {
+  // Create a mock model that simulates relevance checking
+  const createMockModel = (returnValue: boolean): BaseChatModel => {
+    const mockModel = new FakeChatModel({});
+    mockModel.withStructuredOutput = jest.fn().mockReturnValue({
+      pipe: () => ({
+        invoke: async () => ({ answer: returnValue }),
+      }),
+    });
+    (ChatPromptTemplate.fromTemplate as jest.Mock).mockReturnValue({
+      pipe: () => ({
+        invoke: async () => ({ answer: returnValue }),
+      }),
+    });
+    return mockModel;
+  };
+
+  it('should return true for relevant document using string input', async () => {
+    const mockModel = createMockModel(true);
+
+    const result = await checkDocRelevance({
+      doc: 'This is a document about JavaScript programming',
+      query: 'JavaScript development',
+      model: mockModel,
+    });
+    expect(result).toBe(true);
+  });
+
+  it('should return false for irrelevant document using string input', async () => {
+    const mockModel = createMockModel(false);
+    const result = await checkDocRelevance({
+      doc: 'This is a document about cooking recipes',
+      query: 'JavaScript development',
+      model: mockModel,
+    });
+    expect(result).toBe(false);
+  });
+
+  it('should handle Document object input', async () => {
+    const mockModel = createMockModel(true);
+    const doc = new Document({
+      pageContent: 'This is a document about TypeScript',
+      metadata: { id: '1' },
+    });
+    const result = await checkDocRelevance({
+      doc,
+      query: 'TypeScript features',
+      model: mockModel,
+    });
+    expect(result).toBe(true);
+  });
+
+  it('should work with empty document content', async () => {
+    const mockModel = createMockModel(false);
+    const result = await checkDocRelevance({
+      doc: '',
+      query: 'Any query',
+      model: mockModel,
+    });
+    expect(result).toBe(false);
+  });
+
+  it('should work with empty query', async () => {
+    const mockModel = createMockModel(false);
+    const result = await checkDocRelevance({
+      doc: 'Some content',
+      query: '',
+      model: mockModel,
+    });
+    expect(result).toBe(false);
+  });
+
+  it('should handle special characters in document and query', async () => {
+    const mockModel = createMockModel(true);
+    const result = await checkDocRelevance({
+      doc: 'Content with \n newlines and \t tabs',
+      query: 'Query with \n newlines',
+      model: mockModel,
+    });
+    expect(result).toBe(true);
+  });
+});

package/src/ai/utils/__tests__/doc-splitter.test.ts

@@ -0,0 +1,42 @@
+import { Document } from '@langchain/core/documents';
+import { docSplitter } from '../doc-splitter.js';
+
+describe('docSplitter', () => {
+  it('should split a single string into documents', async () => {
+    const text = 'This is a test document';
+    const result = await docSplitter(text);
+    expect(Array.isArray(result)).toBe(true);
+    expect(result.length).toBeGreaterThan(0);
+    expect(result[0]).toBeInstanceOf(Document);
+  });
+
+  it('should split an array of strings into documents', async () => {
+    const texts = ['First document', 'Second document'];
+    const result = await docSplitter(texts);
+    expect(Array.isArray(result)).toBe(true);
+    expect(result.length).toBeGreaterThan(0);
+    expect(result).toEqual(expect.arrayContaining([expect.any(Document)]));
+  });
+
+  it('should throw error for empty string', () => {
+    expect(() => docSplitter('')).toThrow('No text provided');
+  });
+
+  it('should throw error for empty array', () => {
+    expect(() => docSplitter([])).toThrow('Text array cannot be empty');
+  });
+
+  it('should throw error for null or undefined input', () => {
+    // @ts-expect-error testing invalid input
+    expect(() => docSplitter(null)).toThrow('No text provided');
+    // @ts-expect-error testing invalid input
+    expect(() => docSplitter(undefined)).toThrow('No text provided');
+  });
+
+  it('should handle long text correctly', async () => {
+    const longText = 'a'.repeat(2000);
+    const result = await docSplitter(longText);
+    expect(Array.isArray(result)).toBe(true);
+    expect(result.length).toBeGreaterThan(0);
+  });
+});

package/src/ai/utils/__tests__/filter-similarity-search-results.test.ts

@@ -0,0 +1,57 @@
+import { Document } from '@langchain/core/documents';
+import { filterSimilaritySearchResults } from '../filter-similarity-search-results.js';
+
+describe('filterSimilaritySearchResults', () => {
+  const createDoc = (content: string): Document =>
+    new Document({ pageContent: content });
+
+  it('should filter results above threshold', () => {
+    const results = [
+      [createDoc('doc1'), 0.9],
+      [createDoc('doc2'), 0.7],
+      [createDoc('doc3'), 0.3],
+    ] as [Document, number][];
+
+    const filtered = filterSimilaritySearchResults(results, 0.5);
+    expect(filtered).toHaveLength(2);
+    expect(filtered[0]).toBeInstanceOf(Document);
+    expect(filtered[1]).toBeInstanceOf(Document);
+  });
+
+  it('should return empty array when no results meet threshold', () => {
+    const results = [
+      [createDoc('doc1'), 0.3],
+      [createDoc('doc2'), 0.2],
+    ] as [Document, number][];
+
+    const filtered = filterSimilaritySearchResults(results, 0.5);
+    expect(filtered).toHaveLength(0);
+  });
+
+  it('should handle empty results array', () => {
+    const results: [Document, number][] = [];
+    const filtered = filterSimilaritySearchResults(results, 0.5);
+    expect(filtered).toHaveLength(0);
+  });
+
+  it('should handle threshold of 0', () => {
+    const results = [
+      [createDoc('doc1'), 0.1],
+      [createDoc('doc2'), 0.2],
+    ] as [Document, number][];
+
+    const filtered = filterSimilaritySearchResults(results, 0);
+    expect(filtered).toHaveLength(2);
+  });
+
+  it('should handle threshold of 1', () => {
+    const results = [
+      [createDoc('doc1'), 0.9],
+      [createDoc('doc2'), 1.0],
+    ] as [Document, number][];
+
+    const filtered = filterSimilaritySearchResults(results, 1);
+    expect(filtered).toHaveLength(1);
+    expect(filtered[0]).toBeInstanceOf(Document);
+  });
+});

package/src/ai/utils/__tests__/json-to-yaml.test.ts

@@ -0,0 +1,70 @@
+import { jsonToYaml } from '../json-to-yaml.js';
+
+describe('jsonToYaml', () => {
+  it('should convert flat JSON object to YAML', () => {
+    const json = {
+      name: 'John',
+      age: 30,
+      city: 'New York',
+    };
+    const expected = 'name: John\nage: 30\ncity: New York';
+    expect(jsonToYaml(json)).toBe(expected);
+  });
+
+  it('should handle nested objects', () => {
+    const json = {
+      person: {
+        name: 'John',
+        address: {
+          city: 'New York',
+          zip: '10001',
+        },
+      },
+    };
+    const expected =
+      'person:\n  name: John\n  address:\n    city: New York\n    zip: 10001';
+    expect(jsonToYaml(json)).toBe(expected);
+  });
+
+  it('should handle arrays as nested objects', () => {
+    const json = {
+      fruits: ['apple', 'banana', 'orange'],
+    };
+    const expected = 'fruits:\n  0: apple\n  1: banana\n  2: orange';
+    expect(jsonToYaml(json)).toBe(expected);
+  });
+
+  it('should handle empty objects', () => {
+    const json = {};
+    expect(jsonToYaml(json)).toBe('');
+  });
+
+  it('should handle null values', () => {
+    const json = {
+      name: null,
+      age: 30,
+    };
+    const expected = 'name: null\nage: 30';
+    expect(jsonToYaml(json)).toBe(expected);
+  });
+
+  it('should handle mixed types', () => {
+    const json = {
+      string: 'text',
+      number: 42,
+      boolean: true,
+      nullValue: null,
+    };
+    const expected = 'string: text\nnumber: 42\nboolean: true\nnullValue: null';
+    expect(jsonToYaml(json)).toBe(expected);
+  });
+
+  it('should handle custom indentation level', () => {
+    const json = {
+      name: 'John',
+      age: 30,
+    };
+    const expected = '    name: John\n    age: 30';
+    expect(jsonToYaml(json, 2)).toBe(expected);
+  });
+});