Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: rename template to prompt and allow getting history from ragcha… #4

Merged
merged 1 commit into from
May 22, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ export class Config {
public readonly ratelimit?: Ratelimit;

public readonly model?: BaseLanguageModelInterface;
public readonly template?: PromptTemplate;
public readonly prompt?: PromptTemplate;

constructor(config: RAGChatConfig) {
this.vector = config.vector;
Expand All @@ -20,6 +20,6 @@ export class Config {
this.ratelimit = config.ratelimit;

this.model = config.model;
this.template = config.template;
this.prompt = config.prompt;
}
}
2 changes: 1 addition & 1 deletion src/prompts.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import { PromptTemplate } from "@langchain/core/prompts";

export const QA_TEMPLATE =
export const QA_PROMPT_TEMPLATE =
PromptTemplate.fromTemplate(`You are a friendly AI assistant augmented with an Upstash Vector Store.
To help you answer the questions, a context will be provided. This context is generated by querying the vector store with the user question.
Answer the question at the end using only the information available in the context and chat history.
Expand Down
8 changes: 4 additions & 4 deletions src/rag-chat-base.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,18 +17,18 @@ export class RAGChatBase {
protected historyService: HistoryService;

#model: BaseLanguageModelInterface;
#template: PromptTemplate;
#prompt: PromptTemplate;

constructor(
retrievalService: RetrievalService,
historyService: HistoryService,
config: { model: BaseLanguageModelInterface; template: PromptTemplate }
config: { model: BaseLanguageModelInterface; prompt: PromptTemplate }
) {
this.retrievalService = retrievalService;
this.historyService = historyService;

this.#model = config.model;
this.#template = config.template;
this.#prompt = config.prompt;
}

protected async prepareChat({
Expand Down Expand Up @@ -69,7 +69,7 @@ export class RAGChatBase {
question: (input) => input.question,
context: (input) => input.context,
},
this.#template,
this.#prompt,
this.#model,
]);

Expand Down
2 changes: 1 addition & 1 deletion src/rag-chat.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,7 @@ describe("RAG Chat with custom template", () => {
token: process.env.UPSTASH_REDIS_REST_TOKEN!,
url: process.env.UPSTASH_REDIS_REST_URL!,
}),
template: PromptTemplate.fromTemplate("Just say `I'm a cookie monster`. Nothing else."),
prompt: PromptTemplate.fromTemplate("Just say `I'm a cookie monster`. Nothing else."),
model: new ChatOpenAI({
modelName: "gpt-3.5-turbo",
streaming: false,
Expand Down
15 changes: 11 additions & 4 deletions src/rag-chat.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import type { AIMessage } from "@langchain/core/messages";
import type { StreamingTextResponse } from "ai";

import { QA_TEMPLATE } from "./prompts";
import { QA_PROMPT_TEMPLATE } from "./prompts";

import { UpstashModelError } from "./error/model";
import { RatelimitUpstashError } from "./error/ratelimit";
Expand Down Expand Up @@ -29,7 +29,7 @@ export class RAGChat extends RAGChatBase {
}
super(retrievalService, historyService, {
model: config.model,
template: config.template ?? QA_TEMPLATE,
prompt: config.prompt ?? QA_PROMPT_TEMPLATE,
});
this.#ratelimitService = ratelimitService;
}
Expand All @@ -38,7 +38,7 @@ export class RAGChat extends RAGChatBase {
// Adds chat session id and ratelimit session id if not provided.
const options_ = appendDefaultsIfNeeded(options);

//Checks ratelimit of the user. If not enabled `success` will be always true.
// Checks ratelimit of the user. If not enabled `success` will be always true.
const { success, resetTime } = await this.#ratelimitService.checkLimit(
options_.ratelimitSessionId
);
Expand All @@ -50,14 +50,16 @@ export class RAGChat extends RAGChatBase {
});
}

//Sanitizes the given input by stripping all the newline chars then queries vector db with sanitized question.
// Sanitizes the given input by stripping all the newline chars. Then, queries vector db with sanitized question.
const { question, facts } = await this.prepareChat({
question: input,
similarityThreshold: options_.similarityThreshold,
metadataKey: options_.metadataKey,
topK: options_.topK,
});

// Calls LLM service with organized prompt. Prompt holds chat_history, facts gathered from vector db and sanitized question.
// Allows either streaming call via Vercel AI SDK or non-streaming call
return options.stream
? this.streamingChainCall(options_, question, facts)
: this.chainCall(options_, question, facts);
Expand All @@ -71,4 +73,9 @@ export class RAGChat extends RAGChatBase {
);
return retrievalServiceStatus === "Success" ? "OK" : "NOT-OK";
}

/** Method to get history of messages used in the RAG Chat*/
getHistory() {
return this.historyService;
}
}
2 changes: 1 addition & 1 deletion src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ type RAGChatConfigCommon = {
Question: {question}
Helpful answer:`)
*/
template?: PromptTemplate;
prompt?: PromptTemplate;
/**
* Ratelimit instance
* @example new Ratelimit({
Expand Down