rajeshdavidbabu / html-chat-ai-sdk Goto Github PK

import 'server-only';
import { getSession } from '@/app/supabase-server';
import { Database } from '@/lib/db_types';
import { templates } from '@/lib/template';
import { nanoid } from '@/lib/utils';
import { PineconeClient } from '@pinecone-database/pinecone';
import { createServerActionClient } from '@supabase/auth-helpers-nextjs';
import { LangChainStream, Message, StreamingTextResponse } from 'ai';
import { ConversationalRetrievalQAChain, LLMChain } from 'langchain/chains';
import { ChatOpenAI } from 'langchain/chat_models/openai';
import { OpenAIEmbeddings } from 'langchain/embeddings/openai';
import { BufferMemory } from 'langchain/memory';
import { PineconeStore } from 'langchain/vectorstores/pinecone';
import { cookies } from 'next/headers';
import { Configuration, OpenAIApi } from 'openai-edge';

export const runtime = 'nodejs';

const configuration = new Configuration({
  apiKey: process.env.OPENAI_API_KEY
});

const openai = new OpenAIApi(configuration);

const formatMessage = (message: Message) => {
  return `${message.role === 'user' ? 'Human' : 'Assistant'}: ${
    message.content
  }`;
};

export async function POST(req: Request) {
  const cookieStore = cookies();
  const supabase = createServerActionClient<Database>({
    cookies: () => cookieStore
  });
  const session = await getSession();
  const userId = session?.user.id;

  if (!userId) {
    return new Response('Unauthorized', {
      status: 401
    });
  }

  const streamingModel = new ChatOpenAI({
    modelName: 'gpt-3.5-turbo',
    streaming: true,
    verbose: true,
    temperature: 0
  });

  const nonStreamingModel = new ChatOpenAI({
    modelName: 'gpt-3.5-turbo',
    verbose: true,
    temperature: 0
  });
  const pinecone = new PineconeClient();
  await pinecone.init({
    environment: process.env.PINECONE_ENVIRONMENT ?? '',
    apiKey: process.env.PINECONE_API_KEY ?? ''
  });

  const pineconeIndex = pinecone.Index(process.env.PINECONE_INDEX_NAME!);

  const vectorStore = await PineconeStore.fromExistingIndex(
    new OpenAIEmbeddings(),
    { pineconeIndex }
  );

  const json = await req.json();
  const messages: Message[] = json.messages ?? [];
  const formattedPreviousMessages = messages.slice(0, -1).map(formatMessage);
  const question = messages[messages.length - 1].content;
  const sanitizedQuestion = question.trim().replaceAll('\n', ' ');
  const { stream, handlers } = LangChainStream();
  const chatHistory = formattedPreviousMessages.join('\n');

  const chain = ConversationalRetrievalQAChain.fromLLM(
    streamingModel,
    vectorStore.asRetriever(),
    {
      qaTemplate: templates.qaPrompt,
      questionGeneratorTemplate: templates.condensePrompt,
      returnSourceDocuments: true, //default 4
      memory: new BufferMemory({
        memoryKey: 'chat_history',
        inputKey: 'question', // The key for the input to the chain
        outputKey: 'text', // The key for the final conversational output of the chain
        returnMessages: true // If using with a chat model (e.g. gpt-3.5 or gpt-4)
      }),
      questionGeneratorChainOptions: {
        llm: nonStreamingModel
      }
    }
  );

  // Question using chat-history
  // Reference https://js.langchain.com/docs/modules/chains/popular/chat_vector_db#externally-managed-memory
  chain.call(
    {
      question: sanitizedQuestion,
      chat_history: chatHistory
    },
    [handlers]
  )
  return new StreamingTextResponse(stream);
}

And I need to figure out how to implement this logic, expect not using the callback functionality as it breaks the stream and wont save the data to the db 

const { stream, handlers } = LangChainStream({
      async onCompletion(completion) {
        const title = json.messages[0].content.substring(0, 100);
        const id = json.id ?? nanoid();
        const createdAt = Date.now();
        const path = `/chat/${id}`;
        const payload = {
          id,
          title,
          userId,
          createdAt,
          path,
          messages: [
            ...messages,
            {
              content: completion,
              role: 'assistant'
            }
          ]
        };

        await supabase.from('chats').upsert({ id, payload }).throwOnError();
      }
    });