Initial commit

src/agents/imageSearchAgent.ts

@@ -0,0 +1,80 @@
+import {
+  RunnableSequence,
+  RunnableMap,
+  RunnableLambda,
+} from '@langchain/core/runnables';
+import { PromptTemplate } from '@langchain/core/prompts';
+import { OpenAI } from '@langchain/openai';
+import formatChatHistoryAsString from '../utils/formatHistory';
+import { BaseMessage } from '@langchain/core/messages';
+import { StringOutputParser } from '@langchain/core/output_parsers';
+import { searchSearxng } from '../core/searxng';
+
+const llm = new OpenAI({
+  temperature: 0,
+  modelName: 'gpt-3.5-turbo',
+});
+
+const imageSearchChainPrompt = `
+You will be given a conversation below and a follow up question. You need to rephrase the follow-up question so it is a standalone question that can be used by the LLM to search the web for images.
+You need to make sure the rephrased question agrees with the conversation and is relevant to the conversation.
+
+Example:
+1. Follow up question: What is a cat?
+Rephrased: A cat
+
+2. Follow up question: What is a car? How does it works?
+Rephrased: Car working
+
+3. Follow up question: How does an AC work?
+Rephrased: AC working
+
+Conversation:
+{chat_history}
+
+Follow up question: {query}
+Rephrased question:
+`;
+
+type ImageSearchChainInput = {
+  chat_history: BaseMessage[];
+  query: string;
+};
+
+const strParser = new StringOutputParser();
+
+const imageSearchChain = RunnableSequence.from([
+  RunnableMap.from({
+    chat_history: (input: ImageSearchChainInput) => {
+      return formatChatHistoryAsString(input.chat_history);
+    },
+    query: (input: ImageSearchChainInput) => {
+      return input.query;
+    },
+  }),
+  PromptTemplate.fromTemplate(imageSearchChainPrompt),
+  llm,
+  strParser,
+  RunnableLambda.from(async (input: string) => {
+    const res = await searchSearxng(input, {
+      categories: ['images'],
+      engines: ['bing_images', 'google_images'],
+    });
+
+    const images = [];
+
+    res.results.forEach((result) => {
+      if (result.img_src && result.url && result.title) {
+        images.push({
+          img_src: result.img_src,
+          url: result.url,
+          title: result.title,
+        });
+      }
+    });
+
+    return images.slice(0, 10);
+  }),
+]);
+
+export default imageSearchChain;

src/agents/webSearchAgent.ts

@@ -0,0 +1,250 @@
+import { BaseMessage } from '@langchain/core/messages';
+import {
+  PromptTemplate,
+  ChatPromptTemplate,
+  MessagesPlaceholder,
+} from '@langchain/core/prompts';
+import {
+  RunnableSequence,
+  RunnableMap,
+  RunnableLambda,
+} from '@langchain/core/runnables';
+import { ChatOpenAI, OpenAI, OpenAIEmbeddings } from '@langchain/openai';
+import { StringOutputParser } from '@langchain/core/output_parsers';
+import { Document } from '@langchain/core/documents';
+import { searchSearxng } from '../core/searxng';
+import type { StreamEvent } from '@langchain/core/tracers/log_stream';
+import formatChatHistoryAsString from '../utils/formatHistory';
+import eventEmitter from 'events';
+import computeSimilarity from '../utils/computeSimilarity';
+
+const chatLLM = new ChatOpenAI({
+  modelName: 'gpt-3.5-turbo',
+  temperature: 0.7,
+});
+
+const llm = new OpenAI({
+  temperature: 0,
+  modelName: 'gpt-3.5-turbo',
+});
+
+const embeddings = new OpenAIEmbeddings({
+  modelName: 'text-embedding-3-large',
+});
+
+const basicSearchRetrieverPrompt = `
+You will be given a conversation below and a follow up question. You need to rephrase the follow-up question if needed so it is a standalone question that can be used by the LLM to search the web for information.
+If it is a writing task or a simple hi, hello rather than a question, you need to return \`not_needed\` as the response.
+
+Example:
+1. Follow up question: What is the capital of France?
+Rephrased: Capital of france
+
+2. Follow up question: What is the population of New York City?
+Rephrased: Population of New York City
+
+3. Follow up question: What is Docker?
+Rephrased: What is Docker
+
+Conversation:
+{chat_history}
+
+Follow up question: {query}
+Rephrased question:
+`;
+
+const basicWebSearchResponsePrompt = `
+    You are Perplexica, an AI model who is expert at searching the web and answering user's queries.
+
+    Generate a response that is informative and relevant to the user's query based on provided context (the context consits of search results containg a brief description of the content of that page).
+    You must use this context to answer the user's query in the best way possible. Use an unbaised and journalistic tone in your response. Do not repeat the text.
+    You must not tell the user to open any link or visit any website to get the answer. You must provide the answer in the response itself. If the user asks for links you can provide them.
+    Your responses should be medium to long in length be informative and relevant to the user's query. You can use markdowns to format your response. You should use bullet points to list the information. Make sure the answer is not short and is informative.
+    You have to cite the answer using [number] notation. You must cite the sentences with their relevent context number. You must cite each and every part of the answer so the user can know where the information is coming from.
+    Place these citations at the end of that particular sentence. You can cite the same sentence multiple times if it is relevant to the user's query like [number1][number2].
+    However you do not need to cite it using the same number. You can use different numbers to cite the same sentence multiple times. The number refers to the number of the search result (passed in the context) used to generate that part of the answer.
+
+    Aything inside the following \`context\` HTML block provided below is for your knowledge returned by the search engine and is not shared by the user. You have to answer question on the basis of it and cite the relevant information from it but you do not have to 
+    talk about the context in your response. 
+
+    <context>
+    {context}
+    </context>
+
+    If you think there's nothing relevant in the search results, you can say that 'Hmm, sorry I could not find any relevant information on this topic. Would you like me to search again or ask something else?'.
+    Anything between the \`context\` is retrieved from a search engine and is not a part of the conversation with the user. Today's date is ${new Date().toISOString()}
+`;
+
+const strParser = new StringOutputParser();
+
+const handleStream = async (
+  stream: AsyncGenerator<StreamEvent, any, unknown>,
+  emitter: eventEmitter,
+) => {
+  for await (const event of stream) {
+    if (
+      event.event === 'on_chain_end' &&
+      event.name === 'FinalSourceRetriever'
+    ) {
+      emitter.emit(
+        'data',
+        JSON.stringify({ type: 'sources', data: event.data.output }),
+      );
+    }
+    if (
+      event.event === 'on_chain_stream' &&
+      event.name === 'FinalResponseGenerator'
+    ) {
+      emitter.emit(
+        'data',
+        JSON.stringify({ type: 'response', data: event.data.chunk }),
+      );
+    }
+    if (
+      event.event === 'on_chain_end' &&
+      event.name === 'FinalResponseGenerator'
+    ) {
+      emitter.emit('end');
+    }
+  }
+};
+
+const processDocs = async (docs: Document[]) => {
+  return docs
+    .map((_, index) => `${index + 1}. ${docs[index].pageContent}`)
+    .join('\n');
+};
+
+const rerankDocs = async ({
+  query,
+  docs,
+}: {
+  query: string;
+  docs: Document[];
+}) => {
+  if (docs.length === 0) {
+    return docs;
+  }
+
+  const docsWithContent = docs.filter(
+    (doc) => doc.pageContent && doc.pageContent.length > 0,
+  );
+
+  const docEmbeddings = await embeddings.embedDocuments(
+    docsWithContent.map((doc) => doc.pageContent),
+  );
+
+  const queryEmbedding = await embeddings.embedQuery(query);
+
+  const similarity = docEmbeddings.map((docEmbedding, i) => {
+    const sim = computeSimilarity(queryEmbedding, docEmbedding);
+
+    return {
+      index: i,
+      similarity: sim,
+    };
+  });
+
+  const sortedDocs = similarity
+    .sort((a, b) => b.similarity - a.similarity)
+    .filter((sim) => sim.similarity > 0.5)
+    .slice(0, 15)
+    .map((sim) => docsWithContent[sim.index]);
+
+  return sortedDocs;
+};
+
+type BasicChainInput = {
+  chat_history: BaseMessage[];
+  query: string;
+};
+
+const basicWebSearchRetrieverChain = RunnableSequence.from([
+  PromptTemplate.fromTemplate(basicSearchRetrieverPrompt),
+  llm,
+  strParser,
+  RunnableLambda.from(async (input: string) => {
+    if (input === 'not_needed') {
+      return { query: '', docs: [] };
+    }
+
+    const res = await searchSearxng(input, {
+      language: 'en',
+    });
+
+    const documents = res.results.map(
+      (result) =>
+        new Document({
+          pageContent: result.content,
+          metadata: {
+            title: result.title,
+            url: result.url,
+            ...(result.img_src && { img_src: result.img_src }),
+          },
+        }),
+    );
+
+    return { query: input, docs: documents };
+  }),
+]);
+
+const basicWebSearchAnsweringChain = RunnableSequence.from([
+  RunnableMap.from({
+    query: (input: BasicChainInput) => input.query,
+    chat_history: (input: BasicChainInput) => input.chat_history,
+    context: RunnableSequence.from([
+      (input) => ({
+        query: input.query,
+        chat_history: formatChatHistoryAsString(input.chat_history),
+      }),
+      basicWebSearchRetrieverChain
+        .pipe(rerankDocs)
+        .withConfig({
+          runName: 'FinalSourceRetriever',
+        })
+        .pipe(processDocs),
+    ]),
+  }),
+  ChatPromptTemplate.fromMessages([
+    ['system', basicWebSearchResponsePrompt],
+    new MessagesPlaceholder('chat_history'),
+    ['user', '{query}'],
+  ]),
+  chatLLM,
+  strParser,
+]).withConfig({
+  runName: 'FinalResponseGenerator',
+});
+
+const basicWebSearch = (query: string, history: BaseMessage[]) => {
+  const emitter = new eventEmitter();
+
+  try {
+    const stream = basicWebSearchAnsweringChain.streamEvents(
+      {
+        chat_history: history,
+        query: query,
+      },
+      {
+        version: 'v1',
+      },
+    );
+
+    handleStream(stream, emitter);
+  } catch (err) {
+    emitter.emit(
+      'error',
+      JSON.stringify({ data: 'An error has occurred please try again later' }),
+    );
+    console.error(err);
+  }
+
+  return emitter;
+};
+
+const handleWebSearch = (message: string, history: BaseMessage[]) => {
+  const emitter = basicWebSearch(message, history);
+  return emitter;
+};
+
+export default handleWebSearch;

src/app.ts

@@ -0,0 +1,26 @@
+import { startWebSocketServer } from './websocket';
+import express from 'express';
+import cors from 'cors';
+import http from 'http';
+import routes from './routes';
+
+const app = express();
+const server = http.createServer(app);
+
+const corsOptions = {
+  origin: '*',
+};
+
+app.use(cors(corsOptions));
+app.use(express.json());
+
+app.use('/api', routes);
+app.get('/api', (_, res) => {
+  res.status(200).json({ status: 'ok' });
+});
+
+server.listen(process.env.PORT!, () => {
+  console.log(`API server started on port ${process.env.PORT}`);
+});
+
+startWebSocketServer(server);

src/core/agentPicker.ts

@@ -0,0 +1,69 @@
+import { z } from 'zod';
+import { OpenAI } from '@langchain/openai';
+import { RunnableSequence } from '@langchain/core/runnables';
+import { StructuredOutputParser } from 'langchain/output_parsers';
+import { PromptTemplate } from '@langchain/core/prompts';
+
+const availableAgents = [
+  {
+    name: 'webSearch',
+    description:
+      'It is expert is searching the web for information and answer user queries',
+  },
+  /* {
+    name: 'academicSearch',
+    description:
+      'It is expert is searching the academic databases for information and answer user queries. It is particularly good at finding research papers and articles on topics like science, engineering, and technology. Use this instead of wolframAlphaSearch if the user query is not mathematical or scientific in nature',
+  },
+  {
+    name: 'youtubeSearch',
+    description:
+      'This model is expert at finding videos on youtube based on user queries',
+  },
+  {
+    name: 'wolframAlphaSearch',
+    description:
+      'This model is expert at finding answers to mathematical and scientific questions based on user queries.',
+  },
+  {
+    name: 'redditSearch',
+    description:
+      'This model is expert at finding posts and discussions on reddit based on user queries',
+  },
+  {
+    name: 'writingAssistant',
+    description:
+      'If there is no need for searching, this model is expert at generating text based on user queries',
+  }, */
+];
+
+const parser = StructuredOutputParser.fromZodSchema(
+  z.object({
+    agent: z.string().describe('The name of the selected agent'),
+  }),
+);
+
+const prompt = `
+    You are an AI model who is expert at finding suitable agents for user queries. The available agents are:
+    ${availableAgents.map((agent) => `- ${agent.name}: ${agent.description}`).join('\n')}
+
+    Your task is to find the most suitable agent for the following query: {query}
+
+    {format_instructions}
+`;
+
+const chain = RunnableSequence.from([
+  PromptTemplate.fromTemplate(prompt),
+  new OpenAI({ temperature: 0 }),
+  parser,
+]);
+
+const pickSuitableAgent = async (query: string) => {
+  const res = await chain.invoke({
+    query,
+    format_instructions: parser.getFormatInstructions(),
+  });
+  return res.agent;
+};
+
+export default pickSuitableAgent;

src/core/searxng.ts

@@ -0,0 +1,42 @@
+import axios from 'axios';
+
+interface SearxngSearchOptions {
+  categories?: string[];
+  engines?: string[];
+  language?: string;
+  pageno?: number;
+}
+
+interface SearxngSearchResult {
+  title: string;
+  url: string;
+  img_src?: string;
+  thumbnail_src?: string;
+  content?: string;
+  author?: string;
+}
+
+export const searchSearxng = async (
+  query: string,
+  opts?: SearxngSearchOptions,
+) => {
+  const url = new URL(`${process.env.SEARXNG_API_URL}/search?format=json`);
+  url.searchParams.append('q', query);
+
+  if (opts) {
+    Object.keys(opts).forEach((key) => {
+      if (Array.isArray(opts[key])) {
+        url.searchParams.append(key, opts[key].join(','));
+        return;
+      }
+      url.searchParams.append(key, opts[key]);
+    });
+  }
+
+  const res = await axios.get(url.toString());
+
+  const results: SearxngSearchResult[] = res.data.results;
+  const suggestions: string[] = res.data.suggestions;
+
+  return { results, suggestions };
+};

src/routes/images.ts

@@ -0,0 +1,22 @@
+import express from 'express';
+import imageSearchChain from '../agents/imageSearchAgent';
+
+const router = express.Router();
+
+router.post('/', async (req, res) => {
+  try {
+    const { query, chat_history } = req.body;
+
+    const images = await imageSearchChain.invoke({
+      query,
+      chat_history,
+    });
+
+    res.status(200).json({ images });
+  } catch (err) {
+    res.status(500).json({ message: 'An error has occurred.' });
+    console.log(err.message);
+  }
+});
+
+export default router;

src/routes/index.ts

@@ -0,0 +1,8 @@
+import express from 'express';
+import imagesRouter from './images';
+
+const router = express.Router();
+
+router.use('/images', imagesRouter);
+
+export default router;

src/utils/computeSimilarity.ts

@@ -0,0 +1,14 @@
+import dot from 'compute-dot';
+import cosineSimilarity from 'compute-cosine-similarity';
+
+const computeSimilarity = (x: number[], y: number[]): number => {
+  if (process.env.SIMILARITY_MEASURE === 'cosine') {
+    return cosineSimilarity(x, y);
+  } else if (process.env.SIMILARITY_MEASURE === 'dot') {
+    return dot(x, y);
+  }
+
+  throw new Error('Invalid similarity measure');
+};
+
+export default computeSimilarity;

src/utils/formatHistory.ts

@@ -0,0 +1,9 @@
+import { BaseMessage } from '@langchain/core/messages';
+
+const formatChatHistoryAsString = (history: BaseMessage[]) => {
+  return history
+    .map((message) => `${message._getType()}: ${message.content}`)
+    .join('\n');
+};
+
+export default formatChatHistoryAsString;

src/websocket/connectionManager.ts

@@ -0,0 +1,11 @@
+import { WebSocket } from 'ws';
+import { handleMessage } from './messageHandler';
+
+export const handleConnection = (ws: WebSocket) => {
+  ws.on(
+    'message',
+    async (message) => await handleMessage(message.toString(), ws),
+  );
+
+  ws.on('close', () => console.log('Connection closed'));
+};

src/websocket/index.ts

@@ -0,0 +1,8 @@
+import { initServer } from './websocketServer';
+import http from 'http';
+
+export const startWebSocketServer = (
+  server: http.Server<typeof http.IncomingMessage, typeof http.ServerResponse>,
+) => {
+  initServer(server);
+};

src/websocket/messageHandler.ts

@@ -0,0 +1,81 @@
+import { WebSocket } from 'ws';
+import pickSuitableAgent from '../core/agentPicker';
+import handleWebSearch from '../agents/webSearchAgent';
+import { BaseMessage, AIMessage, HumanMessage } from '@langchain/core/messages';
+
+type Message = {
+  type: string;
+  content: string;
+  copilot: boolean;
+  focus: string;
+  history: Array<[string, string]>;
+};
+
+export const handleMessage = async (message: string, ws: WebSocket) => {
+  try {
+    const parsedMessage = JSON.parse(message) as Message;
+    const id = Math.random().toString(36).substring(7);
+
+    if (!parsedMessage.content)
+      return ws.send(
+        JSON.stringify({ type: 'error', data: 'Invalid message format' }),
+      );
+
+    const history: BaseMessage[] = parsedMessage.history.map((msg) => {
+      if (msg[0] === 'human') {
+        return new HumanMessage({
+          content: msg[1],
+        });
+      } else {
+        return new AIMessage({
+          content: msg[1],
+        });
+      }
+    });
+
+    if (parsedMessage.type === 'message') {
+      /* if (!parsedMessage.focus) {
+        const agent = await pickSuitableAgent(parsedMessage.content);
+        parsedMessage.focus = agent;
+      } */
+
+      parsedMessage.focus = 'webSearch';
+
+      switch (parsedMessage.focus) {
+        case 'webSearch': {
+          const emitter = handleWebSearch(parsedMessage.content, history);
+          emitter.on('data', (data) => {
+            const parsedData = JSON.parse(data);
+            if (parsedData.type === 'response') {
+              ws.send(
+                JSON.stringify({
+                  type: 'message',
+                  data: parsedData.data,
+                  messageId: id,
+                }),
+              );
+            } else if (parsedData.type === 'sources') {
+              ws.send(
+                JSON.stringify({
+                  type: 'sources',
+                  data: parsedData.data,
+                  messageId: id,
+                }),
+              );
+            }
+          });
+          emitter.on('end', () => {
+            ws.send(JSON.stringify({ type: 'messageEnd', messageId: id }));
+          });
+          emitter.on('error', (data) => {
+            const parsedData = JSON.parse(data);
+            ws.send(JSON.stringify({ type: 'error', data: parsedData.data }));
+          });
+        }
+      }
+    }
+  } catch (error) {
+    console.error('Failed to handle message', error);
+    ws.send(JSON.stringify({ type: 'error', data: 'Invalid message format' }));
+  }
+};

src/websocket/websocketServer.ts

@@ -0,0 +1,15 @@
+import { WebSocketServer } from 'ws';
+import { handleConnection } from './connectionManager';
+import http from 'http';
+
+export const initServer = (
+  server: http.Server<typeof http.IncomingMessage, typeof http.ServerResponse>,
+) => {
+  const wss = new WebSocketServer({ server });
+
+  wss.on('connection', (ws) => {
+    handleConnection(ws);
+  });
+
+  console.log(`WebSocket server started on port ${process.env.PORT}`);
+};